diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 3494ae7..8dc1a21 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -5072,13 +5072,13 @@
                                      operation->ctx.tls12_prf.label_length);
             mbedtls_free(operation->ctx.tls12_prf.label);
         }
-
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_TLS12_PSK_TO_MS)
         if (operation->ctx.tls12_prf.other_secret != NULL) {
             mbedtls_platform_zeroize(operation->ctx.tls12_prf.other_secret,
                                      operation->ctx.tls12_prf.other_secret_length);
             mbedtls_free(operation->ctx.tls12_prf.other_secret);
         }
-
+#endif /* MBEDTLS_PSA_BUILTIN_ALG_TLS12_PSK_TO_MS */
         status = PSA_SUCCESS;
 
         /* We leave the fields Ai and output_block to be erased safely by the
@@ -7163,24 +7163,10 @@
     return status;
 }
 
-#if defined(MBEDTLS_PSA_BUILTIN_PAKE)
 psa_status_t psa_pake_setup(
     psa_pake_operation_t *operation,
     const psa_pake_cipher_suite_t *cipher_suite)
 {
-    /* A context must be freshly initialized before it can be set up. */
-    if (operation->alg != PSA_ALG_NONE) {
-        return PSA_ERROR_BAD_STATE;
-    }
-
-    if (cipher_suite == NULL ||
-        PSA_ALG_IS_PAKE(cipher_suite->algorithm) == 0 ||
-        (cipher_suite->type != PSA_PAKE_PRIMITIVE_TYPE_ECC &&
-         cipher_suite->type != PSA_PAKE_PRIMITIVE_TYPE_DH) ||
-        PSA_ALG_IS_HASH(cipher_suite->hash) == 0) {
-        return PSA_ERROR_INVALID_ARGUMENT;
-    }
-
     return psa_driver_wrapper_pake_setup(operation, cipher_suite);
 }
 
@@ -7189,34 +7175,30 @@
     mbedtls_svc_key_id_t password)
 {
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
-    psa_key_attributes_t attributes = psa_key_attributes_init();
-    psa_key_type_t type;
-    psa_key_usage_t usage;
+    psa_status_t unlock_status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_key_slot_t *slot = NULL;
 
-    if (operation->alg == PSA_ALG_NONE) {
+    if (operation->id == 0) {
         return PSA_ERROR_BAD_STATE;
     }
 
-    status = psa_get_key_attributes(password, &attributes);
+    status = psa_get_and_lock_key_slot_with_policy(password, &slot,
+                                                   PSA_KEY_USAGE_DERIVE,
+                                                   PSA_ALG_JPAKE);
     if (status != PSA_SUCCESS) {
         return status;
     }
 
-    type = psa_get_key_type(&attributes);
-    usage = psa_get_key_usage_flags(&attributes);
+    psa_key_attributes_t attributes = {
+        .core = slot->attr
+    };
 
-    psa_reset_key_attributes(&attributes);
+    status = psa_driver_wrapper_pake_set_password_key(&attributes, operation,
+                                                      slot->key.data, slot->key.bytes);
 
-    if (type != PSA_KEY_TYPE_PASSWORD &&
-        type != PSA_KEY_TYPE_PASSWORD_HASH) {
-        return PSA_ERROR_INVALID_ARGUMENT;
-    }
+    unlock_status = psa_unlock_key_slot(slot);
 
-    if ((usage & PSA_KEY_USAGE_DERIVE) == 0) {
-        return PSA_ERROR_NOT_PERMITTED;
-    }
-
-    return psa_driver_wrapper_pake_set_password_key(operation, password);
+    return (status == PSA_SUCCESS) ? unlock_status : status;
 }
 
 psa_status_t psa_pake_set_user(
@@ -7224,7 +7206,7 @@
     const uint8_t *user_id,
     size_t user_id_len)
 {
-    if (operation->alg == PSA_ALG_NONE) {
+    if (operation->id == 0) {
         return PSA_ERROR_BAD_STATE;
     }
 
@@ -7241,7 +7223,7 @@
     const uint8_t *peer_id,
     size_t peer_id_len)
 {
-    if (operation->alg == PSA_ALG_NONE) {
+    if (operation->id == 0) {
         return PSA_ERROR_BAD_STATE;
     }
 
@@ -7257,7 +7239,7 @@
     psa_pake_operation_t *operation,
     psa_pake_role_t role)
 {
-    if (operation->alg == PSA_ALG_NONE) {
+    if (operation->id == 0) {
         return PSA_ERROR_BAD_STATE;
     }
 
@@ -7279,7 +7261,7 @@
     size_t output_size,
     size_t *output_length)
 {
-    if (operation->alg == PSA_ALG_NONE) {
+    if (operation->id == 0) {
         return PSA_ERROR_BAD_STATE;
     }
 
@@ -7297,7 +7279,7 @@
     const uint8_t *input,
     size_t input_length)
 {
-    if (operation->alg == PSA_ALG_NONE) {
+    if (operation->id == 0) {
         return PSA_ERROR_BAD_STATE;
     }
 
@@ -7313,7 +7295,7 @@
     psa_pake_operation_t *operation,
     psa_key_derivation_operation_t *output)
 {
-    if (operation->alg == PSA_ALG_NONE) {
+    if (operation->id == 0) {
         return PSA_ERROR_BAD_STATE;
     }
 
@@ -7323,12 +7305,12 @@
 psa_status_t psa_pake_abort(
     psa_pake_operation_t *operation)
 {
-    if (operation->alg == PSA_ALG_NONE) {
+    /* Aborting a non-active operation is allowed */
+    if (operation->id == 0) {
         return PSA_SUCCESS;
     }
 
     return psa_driver_wrapper_pake_abort(operation);
 }
-#endif /* MBEDTLS_PSA_BUILTIN_PAKE */
 
 #endif /* MBEDTLS_PSA_CRYPTO_C */
diff --git a/library/psa_crypto_pake.c b/library/psa_crypto_pake.c
index 0dafe78..6c4db6f 100644
--- a/library/psa_crypto_pake.c
+++ b/library/psa_crypto_pake.c
@@ -191,9 +191,26 @@
 #endif
 
 #if defined(MBEDTLS_PSA_BUILTIN_PAKE)
-psa_status_t mbedtls_psa_pake_setup(psa_pake_operation_t *operation,
+psa_status_t mbedtls_psa_pake_setup(mbedtls_psa_pake_operation_t *operation,
                                     const psa_pake_cipher_suite_t *cipher_suite)
 {
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+
+    /* A context must be freshly initialized before it can be set up. */
+    if (operation->alg != PSA_ALG_NONE) {
+        status = PSA_ERROR_BAD_STATE;
+        goto error;
+    }
+
+    if (cipher_suite == NULL ||
+        PSA_ALG_IS_PAKE(cipher_suite->algorithm) == 0 ||
+        (cipher_suite->type != PSA_PAKE_PRIMITIVE_TYPE_ECC &&
+         cipher_suite->type != PSA_PAKE_PRIMITIVE_TYPE_DH) ||
+        PSA_ALG_IS_HASH(cipher_suite->hash) == 0) {
+        status = PSA_ERROR_INVALID_ARGUMENT;
+        goto error;
+    }
+
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
     if (cipher_suite->algorithm == PSA_ALG_JPAKE) {
         if (cipher_suite->type != PSA_PAKE_PRIMITIVE_TYPE_ECC ||
@@ -206,12 +223,14 @@
 
         operation->alg = cipher_suite->algorithm;
 
-        mbedtls_ecjpake_init(&operation->ctx.ecjpake);
+        mbedtls_ecjpake_init(&operation->ctx.pake);
 
         operation->state = PSA_PAKE_STATE_SETUP;
         operation->sequence = PSA_PAKE_SEQ_INVALID;
         operation->input_step = PSA_PAKE_STEP_X1_X2;
         operation->output_step = PSA_PAKE_STEP_X1_X2;
+        operation->password_len = 0;
+        operation->password = NULL;
 
         mbedtls_platform_zeroize(operation->buffer, MBEDTLS_PSA_PAKE_BUFFER_SIZE);
         operation->buffer_length = 0;
@@ -230,30 +249,14 @@
     return status;
 }
 
-psa_status_t mbedtls_psa_pake_set_password_key(psa_pake_operation_t *operation,
-                                               mbedtls_svc_key_id_t password)
+psa_status_t mbedtls_psa_pake_set_password_key(const psa_key_attributes_t *attributes,
+                                               mbedtls_psa_pake_operation_t *operation,
+                                               uint8_t *password,
+                                               size_t password_len)
 {
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
-    psa_key_attributes_t attributes = psa_key_attributes_init();
-    psa_key_type_t type;
-    psa_key_usage_t usage;
-    psa_key_slot_t *slot = NULL;
-
-    if (operation->alg == PSA_ALG_NONE ||
-        operation->state != PSA_PAKE_STATE_SETUP) {
-        status = PSA_ERROR_BAD_STATE;
-        goto error;
-    }
-
-    status = psa_get_key_attributes(password, &attributes);
-    if (status != PSA_SUCCESS) {
-        goto error;
-    }
-
-    type = psa_get_key_type(&attributes);
-    usage = psa_get_key_usage_flags(&attributes);
-
-    psa_reset_key_attributes(&attributes);
+    psa_key_type_t type = psa_get_key_type(attributes);
+    psa_key_usage_t usage = psa_get_key_usage_flags(attributes);
 
     if (type != PSA_KEY_TYPE_PASSWORD &&
         type != PSA_KEY_TYPE_PASSWORD_HASH) {
@@ -266,44 +269,48 @@
         goto error;
     }
 
+    if (operation->alg == PSA_ALG_NONE) {
+        status = PSA_ERROR_BAD_STATE;
+        goto error;
+    }
+
+    if (operation->state != PSA_PAKE_STATE_SETUP) {
+        status =  PSA_ERROR_BAD_STATE;
+        goto error;
+    }
+
     if (operation->password != NULL) {
-        return PSA_ERROR_BAD_STATE;
+        status = PSA_ERROR_BAD_STATE;
+        goto error;
     }
 
-    status = psa_get_and_lock_key_slot_with_policy(password, &slot,
-                                                   PSA_KEY_USAGE_DERIVE,
-                                                   PSA_ALG_JPAKE);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    operation->password = mbedtls_calloc(1, slot->key.bytes);
+    operation->password = mbedtls_calloc(1, password_len);
     if (operation->password == NULL) {
-        psa_unlock_key_slot(slot);
         return PSA_ERROR_INSUFFICIENT_MEMORY;
     }
-    memcpy(operation->password, slot->key.data, slot->key.bytes);
-    operation->password_len = slot->key.bytes;
 
-    status = psa_unlock_key_slot(slot);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
+    memcpy(operation->password, password, password_len);
+    operation->password_len = password_len;
 
     return PSA_SUCCESS;
 
 error:
-    psa_pake_abort(operation);
+    mbedtls_psa_pake_abort(operation);
     return status;
 }
 
-psa_status_t mbedtls_psa_pake_set_user(psa_pake_operation_t *operation,
+psa_status_t mbedtls_psa_pake_set_user(mbedtls_psa_pake_operation_t *operation,
                                        const uint8_t *user_id,
                                        size_t user_id_len)
 {
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
     (void) user_id;
     (void) user_id_len;
 
+    if (operation->alg == PSA_ALG_NONE) {
+        return PSA_ERROR_BAD_STATE;
+    }
+
     if (operation->state != PSA_PAKE_STATE_SETUP) {
         status = PSA_ERROR_BAD_STATE;
         goto error;
@@ -316,13 +323,19 @@
     return status;
 }
 
-psa_status_t mbedtls_psa_pake_set_peer(psa_pake_operation_t *operation,
+psa_status_t mbedtls_psa_pake_set_peer(mbedtls_psa_pake_operation_t *operation,
                                        const uint8_t *peer_id,
                                        size_t peer_id_len)
 {
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
     (void) peer_id;
     (void) peer_id_len;
 
+    if (operation->alg == PSA_ALG_NONE) {
+        status =  PSA_ERROR_BAD_STATE;
+        goto error;
+    }
+
     if (operation->state != PSA_PAKE_STATE_SETUP) {
         status = PSA_ERROR_BAD_STATE;
         goto error;
@@ -335,9 +348,15 @@
     return status;
 }
 
-psa_status_t mbedtls_psa_pake_set_role(psa_pake_operation_t *operation,
+psa_status_t mbedtls_psa_pake_set_role(mbedtls_psa_pake_operation_t *operation,
                                        psa_pake_role_t role)
 {
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    if (operation->alg == PSA_ALG_NONE) {
+        status = PSA_ERROR_BAD_STATE;
+        goto error;
+    }
+
     if (operation->state != PSA_PAKE_STATE_SETUP) {
         status = PSA_ERROR_BAD_STATE;
         goto error;
@@ -366,7 +385,7 @@
 }
 
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
-static psa_status_t psa_pake_ecjpake_setup(psa_pake_operation_t *operation)
+static psa_status_t psa_pake_ecjpake_setup(mbedtls_psa_pake_operation_t *operation)
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
     mbedtls_ecjpake_role role;
@@ -383,7 +402,7 @@
         return PSA_ERROR_BAD_STATE;
     }
 
-    ret = mbedtls_ecjpake_setup(&operation->ctx.ecjpake,
+    ret = mbedtls_ecjpake_setup(&operation->ctx.pake,
                                 role,
                                 MBEDTLS_MD_SHA256,
                                 MBEDTLS_ECP_DP_SECP256R1,
@@ -406,7 +425,7 @@
 #endif
 
 static psa_status_t mbedtls_psa_pake_output_internal(
-    psa_pake_operation_t *operation,
+    mbedtls_psa_pake_operation_t *operation,
     psa_pake_step_t step,
     uint8_t *output,
     size_t output_size,
@@ -416,6 +435,10 @@
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
     size_t length;
 
+    if (operation->alg == PSA_ALG_NONE) {
+        return PSA_ERROR_BAD_STATE;
+    }
+
     if (operation->state == PSA_PAKE_STATE_INVALID) {
         return PSA_ERROR_BAD_STATE;
     }
@@ -504,7 +527,7 @@
         /* Initialize & write round on KEY_SHARE sequences */
         if (operation->state == PSA_PAKE_OUTPUT_X1_X2 &&
             operation->sequence == PSA_PAKE_X1_STEP_KEY_SHARE) {
-            ret = mbedtls_ecjpake_write_round_one(&operation->ctx.ecjpake,
+            ret = mbedtls_ecjpake_write_round_one(&operation->ctx.pake,
                                                   operation->buffer,
                                                   MBEDTLS_PSA_PAKE_BUFFER_SIZE,
                                                   &operation->buffer_length,
@@ -517,7 +540,7 @@
             operation->buffer_offset = 0;
         } else if (operation->state == PSA_PAKE_OUTPUT_X2S &&
                    operation->sequence == PSA_PAKE_X1_STEP_KEY_SHARE) {
-            ret = mbedtls_ecjpake_write_round_two(&operation->ctx.ecjpake,
+            ret = mbedtls_ecjpake_write_round_two(&operation->ctx.pake,
                                                   operation->buffer,
                                                   MBEDTLS_PSA_PAKE_BUFFER_SIZE,
                                                   &operation->buffer_length,
@@ -594,7 +617,7 @@
     { return PSA_ERROR_NOT_SUPPORTED; }
 }
 
-psa_status_t mbedtls_psa_pake_output(psa_pake_operation_t *operation,
+psa_status_t mbedtls_psa_pake_output(mbedtls_psa_pake_operation_t *operation,
                                      psa_pake_step_t step,
                                      uint8_t *output,
                                      size_t output_size,
@@ -604,14 +627,14 @@
         operation, step, output, output_size, output_length);
 
     if (status != PSA_SUCCESS) {
-        psa_pake_abort(operation);
+        mbedtls_psa_pake_abort(operation);
     }
 
     return status;
 }
 
 static psa_status_t mbedtls_psa_pake_input_internal(
-    psa_pake_operation_t *operation,
+    mbedtls_psa_pake_operation_t *operation,
     psa_pake_step_t step,
     const uint8_t *input,
     size_t input_length)
@@ -619,6 +642,10 @@
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
 
+    if (operation->alg == PSA_ALG_NONE) {
+        return PSA_ERROR_BAD_STATE;
+    }
+
     if (operation->state == PSA_PAKE_STATE_INVALID) {
         return PSA_ERROR_BAD_STATE;
     }
@@ -746,7 +773,7 @@
         /* Load buffer at each last round ZK_PROOF */
         if (operation->state == PSA_PAKE_INPUT_X1_X2 &&
             operation->sequence == PSA_PAKE_X2_STEP_ZK_PROOF) {
-            ret = mbedtls_ecjpake_read_round_one(&operation->ctx.ecjpake,
+            ret = mbedtls_ecjpake_read_round_one(&operation->ctx.pake,
                                                  operation->buffer,
                                                  operation->buffer_length);
 
@@ -758,7 +785,7 @@
             }
         } else if (operation->state == PSA_PAKE_INPUT_X4S &&
                    operation->sequence == PSA_PAKE_X1_STEP_ZK_PROOF) {
-            ret = mbedtls_ecjpake_read_round_two(&operation->ctx.ecjpake,
+            ret = mbedtls_ecjpake_read_round_two(&operation->ctx.pake,
                                                  operation->buffer,
                                                  operation->buffer_length);
 
@@ -791,7 +818,7 @@
     { return PSA_ERROR_NOT_SUPPORTED; }
 }
 
-psa_status_t mbedtls_psa_pake_input(psa_pake_operation_t *operation,
+psa_status_t mbedtls_psa_pake_input(mbedtls_psa_pake_operation_t *operation,
                                     psa_pake_step_t step,
                                     const uint8_t *input,
                                     size_t input_length)
@@ -800,19 +827,23 @@
         operation, step, input, input_length);
 
     if (status != PSA_SUCCESS) {
-        psa_pake_abort(operation);
+        mbedtls_psa_pake_abort(operation);
     }
 
     return status;
 }
 
 psa_status_t mbedtls_psa_pake_get_implicit_key(
-    psa_pake_operation_t *operation,
+    mbedtls_psa_pake_operation_t *operation,
     psa_key_derivation_operation_t *output)
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
 
+    if (operation->alg == PSA_ALG_NONE) {
+        return PSA_ERROR_BAD_STATE;
+    }
+
     if (operation->input_step != PSA_PAKE_STEP_DERIVE ||
         operation->output_step != PSA_PAKE_STEP_DERIVE) {
         status = PSA_ERROR_BAD_STATE;
@@ -821,14 +852,14 @@
 
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
     if (operation->alg == PSA_ALG_JPAKE) {
-        ret = mbedtls_ecjpake_write_shared_key(&operation->ctx.ecjpake,
+        ret = mbedtls_ecjpake_write_shared_key(&operation->ctx.pake,
                                                operation->buffer,
                                                MBEDTLS_PSA_PAKE_BUFFER_SIZE,
                                                &operation->buffer_length,
                                                mbedtls_psa_get_random,
                                                MBEDTLS_PSA_RANDOM_STATE);
         if (ret != 0) {
-            psa_pake_abort(operation);
+            mbedtls_psa_pake_abort(operation);
             return mbedtls_ecjpake_to_psa_error(ret);
         }
 
@@ -839,7 +870,7 @@
 
         mbedtls_platform_zeroize(operation->buffer, MBEDTLS_PSA_PAKE_BUFFER_SIZE);
 
-        psa_pake_abort(operation);
+        mbedtls_psa_pake_abort(operation);
 
         return status;
     } else
@@ -850,14 +881,19 @@
 
 error:
     psa_key_derivation_abort(output);
-    psa_pake_abort(operation);
+    mbedtls_psa_pake_abort(operation);
 
     return status;
 }
 
-psa_status_t mbedtls_psa_pake_abort(psa_pake_operation_t *operation)
+psa_status_t mbedtls_psa_pake_abort(mbedtls_psa_pake_operation_t *operation)
 {
+    if (operation->alg == PSA_ALG_NONE) {
+        return PSA_SUCCESS;
+    }
+
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
+
     if (operation->alg == PSA_ALG_JPAKE) {
         operation->input_step = PSA_PAKE_STEP_INVALID;
         operation->output_step = PSA_PAKE_STEP_INVALID;
@@ -871,7 +907,7 @@
         mbedtls_platform_zeroize(operation->buffer, MBEDTLS_PSA_PAKE_BUFFER_SIZE);
         operation->buffer_length = 0;
         operation->buffer_offset = 0;
-        mbedtls_ecjpake_free(&operation->ctx.ecjpake);
+        mbedtls_ecjpake_free(&operation->ctx.pake);
     }
 #endif
 
diff --git a/library/psa_crypto_pake.h b/library/psa_crypto_pake.h
index b61ddde..c7bf270 100644
--- a/library/psa_crypto_pake.h
+++ b/library/psa_crypto_pake.h
@@ -93,14 +93,15 @@
  *         It is implementation-dependent whether a failure to initialize
  *         results in this error code.
  */
-psa_status_t mbedtls_psa_pake_setup(psa_pake_operation_t *operation,
+psa_status_t mbedtls_psa_pake_setup(mbedtls_psa_pake_operation_t *operation,
                                     const psa_pake_cipher_suite_t *cipher_suite);
 
 /** Set the password for a password-authenticated key exchange from key ID.
  *
  * Call this function when the password, or a value derived from the password,
  * is already present in the key store.
- *
+ * \param[in] attributes        The attributes of the key to use for the
+ *                              operation.
  * \param[in,out] operation     The operation object to set the password for. It
  *                              must have been set up by psa_pake_setup() and
  *                              not yet in use (neither psa_pake_output() nor
@@ -108,13 +109,8 @@
  *                              be on operation for which the password hasn't
  *                              been set yet (psa_pake_set_password_key()
  *                              hasn't been called yet).
- * \param password              Identifier of the key holding the password or a
- *                              value derived from the password (eg. by a
- *                              memory-hard function).  It must remain valid
- *                              until the operation terminates. It must be of
- *                              type #PSA_KEY_TYPE_PASSWORD or
- *                              #PSA_KEY_TYPE_PASSWORD_HASH. It has to allow
- *                              the usage #PSA_KEY_USAGE_DERIVE.
+ * \param password              Buffer holding the password
+ * \param password_len          Password buffer size
  *
  * \retval #PSA_SUCCESS
  *         Success.
@@ -142,8 +138,10 @@
  *         results in this error code.
  */
 psa_status_t mbedtls_psa_pake_set_password_key(
-    psa_pake_operation_t *operation,
-    mbedtls_svc_key_id_t password);
+    const psa_key_attributes_t *attributes,
+    mbedtls_psa_pake_operation_t *operation,
+    uint8_t *password,
+    size_t password_len);
 
 /** Set the user ID for a password-authenticated key exchange.
  *
@@ -182,7 +180,7 @@
  *         It is implementation-dependent whether a failure to initialize
  *         results in this error code.
  */
-psa_status_t mbedtls_psa_pake_set_user(psa_pake_operation_t *operation,
+psa_status_t mbedtls_psa_pake_set_user(mbedtls_psa_pake_operation_t *operation,
                                        const uint8_t *user_id,
                                        size_t user_id_len);
 
@@ -224,7 +222,7 @@
  *         It is implementation-dependent whether a failure to initialize
  *         results in this error code.
  */
-psa_status_t mbedtls_psa_pake_set_peer(psa_pake_operation_t *operation,
+psa_status_t mbedtls_psa_pake_set_peer(mbedtls_psa_pake_operation_t *operation,
                                        const uint8_t *peer_id,
                                        size_t peer_id_len);
 
@@ -266,7 +264,7 @@
  *         It is implementation-dependent whether a failure to initialize
  *         results in this error code.
  */
-psa_status_t mbedtls_psa_pake_set_role(psa_pake_operation_t *operation,
+psa_status_t mbedtls_psa_pake_set_role(mbedtls_psa_pake_operation_t *operation,
                                        psa_pake_role_t role);
 
 /** Get output for a step of a password-authenticated key exchange.
@@ -324,7 +322,7 @@
  *         It is implementation-dependent whether a failure to initialize
  *         results in this error code.
  */
-psa_status_t mbedtls_psa_pake_output(psa_pake_operation_t *operation,
+psa_status_t mbedtls_psa_pake_output(mbedtls_psa_pake_operation_t *operation,
                                      psa_pake_step_t step,
                                      uint8_t *output,
                                      size_t output_size,
@@ -379,7 +377,7 @@
  *         It is implementation-dependent whether a failure to initialize
  *         results in this error code.
  */
-psa_status_t mbedtls_psa_pake_input(psa_pake_operation_t *operation,
+psa_status_t mbedtls_psa_pake_input(mbedtls_psa_pake_operation_t *operation,
                                     psa_pake_step_t step,
                                     const uint8_t *input,
                                     size_t input_length);
@@ -443,7 +441,7 @@
  *         results in this error code.
  */
 psa_status_t mbedtls_psa_pake_get_implicit_key(
-    psa_pake_operation_t *operation,
+    mbedtls_psa_pake_operation_t *operation,
     psa_key_derivation_operation_t *output);
 
 /** Abort a PAKE operation.
@@ -470,6 +468,6 @@
  *         It is implementation-dependent whether a failure to initialize
  *         results in this error code.
  */
-psa_status_t mbedtls_psa_pake_abort(psa_pake_operation_t *operation);
+psa_status_t mbedtls_psa_pake_abort(mbedtls_psa_pake_operation_t *operation);
 
 #endif /* PSA_CRYPTO_PAKE_H */