diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 18aa18b..4742c3c 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -7180,7 +7180,29 @@
     psa_pake_operation_t *operation,
     const psa_pake_cipher_suite_t *cipher_suite)
 {
-    return psa_driver_wrapper_pake_setup(operation, cipher_suite);
+    if (operation->stage != PSA_PAKE_OPERATION_STAGE_COLLECT_INPUTS) {
+        return PSA_ERROR_BAD_STATE;
+    }
+
+    if (operation->data.inputs.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;
+    }
+
+    ;
+    memset(&operation->data.inputs, 0, sizeof(operation->data.inputs));
+
+    operation->data.inputs.alg = cipher_suite->algorithm;
+    operation->data.inputs.cipher_suite = *cipher_suite;
+
+    return PSA_SUCCESS;
 }
 
 psa_status_t psa_pake_set_password_key(
@@ -7191,7 +7213,11 @@
     psa_status_t unlock_status = PSA_ERROR_CORRUPTION_DETECTED;
     psa_key_slot_t *slot = NULL;
 
-    if (operation->id == 0) {
+    if (operation->stage != PSA_PAKE_OPERATION_STAGE_COLLECT_INPUTS) {
+        return PSA_ERROR_BAD_STATE;
+    }
+
+    if (operation->data.inputs.alg == PSA_ALG_NONE) {
         return PSA_ERROR_BAD_STATE;
     }
 
@@ -7206,9 +7232,29 @@
         .core = slot->attr
     };
 
-    status = psa_driver_wrapper_pake_set_password_key(&attributes, operation,
-                                                      slot->key.data, slot->key.bytes);
+    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) {
+        status = PSA_ERROR_INVALID_ARGUMENT;
+        goto error;
+    }
+
+    if ((usage & PSA_KEY_USAGE_DERIVE) == 0) {
+        status = PSA_ERROR_NOT_PERMITTED;
+        goto error;
+    }
+
+    operation->data.inputs.password = mbedtls_calloc(1, slot->key.bytes);
+    if (operation->data.inputs.password == NULL) {
+        return PSA_ERROR_INSUFFICIENT_MEMORY;
+    }
+
+    memcpy(operation->data.inputs.password, slot->key.data, slot->key.bytes);
+    operation->data.inputs.password_len = slot->key.bytes;
+    operation->data.inputs.key_lifetime = attributes.core.lifetime;
+error:
     unlock_status = psa_unlock_key_slot(slot);
 
     return (status == PSA_SUCCESS) ? unlock_status : status;
@@ -7219,16 +7265,21 @@
     const uint8_t *user_id,
     size_t user_id_len)
 {
-    if (operation->id == 0) {
+    (void) user_id;
+
+    if (operation->stage != PSA_PAKE_OPERATION_STAGE_COLLECT_INPUTS) {
         return PSA_ERROR_BAD_STATE;
     }
 
-    if (user_id_len == 0 || user_id == NULL) {
+    if (operation->data.inputs.alg == PSA_ALG_NONE) {
+        return PSA_ERROR_BAD_STATE;
+    }
+
+    if (user_id_len == 0) {
         return PSA_ERROR_INVALID_ARGUMENT;
     }
 
-    return psa_driver_wrapper_pake_set_user(operation, user_id,
-                                            user_id_len);
+    return PSA_ERROR_NOT_SUPPORTED;
 }
 
 psa_status_t psa_pake_set_peer(
@@ -7236,23 +7287,32 @@
     const uint8_t *peer_id,
     size_t peer_id_len)
 {
-    if (operation->id == 0) {
+    (void) peer_id;
+
+    if (operation->stage != PSA_PAKE_OPERATION_STAGE_COLLECT_INPUTS) {
         return PSA_ERROR_BAD_STATE;
     }
 
-    if (peer_id_len == 0 || peer_id == NULL) {
+    if (operation->data.inputs.alg == PSA_ALG_NONE) {
+        return PSA_ERROR_BAD_STATE;
+    }
+
+    if (peer_id_len == 0) {
         return PSA_ERROR_INVALID_ARGUMENT;
     }
 
-    return psa_driver_wrapper_pake_set_peer(operation, peer_id,
-                                            peer_id_len);
+    return PSA_ERROR_NOT_SUPPORTED;
 }
 
 psa_status_t psa_pake_set_role(
     psa_pake_operation_t *operation,
     psa_pake_role_t role)
 {
-    if (operation->id == 0) {
+    if (operation->stage != PSA_PAKE_OPERATION_STAGE_COLLECT_INPUTS) {
+        return PSA_ERROR_BAD_STATE;
+    }
+
+    if (operation->data.inputs.alg == PSA_ALG_NONE) {
         return PSA_ERROR_BAD_STATE;
     }
 
@@ -7264,7 +7324,9 @@
         return PSA_ERROR_INVALID_ARGUMENT;
     }
 
-    return psa_driver_wrapper_pake_set_role(operation, role);
+    operation->data.inputs.role = role;
+
+    return PSA_SUCCESS;
 }
 
 psa_status_t psa_pake_output(
@@ -7274,11 +7336,34 @@
     size_t output_size,
     size_t *output_length)
 {
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+
+    if (operation->stage == PSA_PAKE_OPERATION_STAGE_COLLECT_INPUTS) {
+        if (operation->data.inputs.alg == PSA_ALG_NONE ||
+            operation->data.inputs.password_len == 0 ||
+            operation->data.inputs.role == PSA_PAKE_ROLE_NONE) {
+            return PSA_ERROR_BAD_STATE;
+        }
+
+        status = psa_driver_wrapper_pake_setup(operation,
+                                               &operation->data.inputs);
+
+        if (status == PSA_SUCCESS) {
+            operation->stage = PSA_PAKE_OPERATION_STAGE_COMPUTATION;
+        } else {
+            return status;
+        }
+    }
+
+    if (operation->stage != PSA_PAKE_OPERATION_STAGE_COMPUTATION) {
+        return PSA_ERROR_BAD_STATE;
+    }
+
     if (operation->id == 0) {
         return PSA_ERROR_BAD_STATE;
     }
 
-    if (output == NULL || output_size == 0 || output_length == NULL) {
+    if (output == NULL || output_size == 0) {
         return PSA_ERROR_INVALID_ARGUMENT;
     }
 
@@ -7292,6 +7377,29 @@
     const uint8_t *input,
     size_t input_length)
 {
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+
+    if (operation->stage == PSA_PAKE_OPERATION_STAGE_COLLECT_INPUTS) {
+        if (operation->data.inputs.alg == PSA_ALG_NONE ||
+            operation->data.inputs.password_len == 0 ||
+            operation->data.inputs.role == PSA_PAKE_ROLE_NONE) {
+            return PSA_ERROR_BAD_STATE;
+        }
+
+        status = psa_driver_wrapper_pake_setup(operation,
+                                               &operation->data.inputs);
+
+        if (status == PSA_SUCCESS) {
+            operation->stage = PSA_PAKE_OPERATION_STAGE_COMPUTATION;
+        } else {
+            return status;
+        }
+    }
+
+    if (operation->stage != PSA_PAKE_OPERATION_STAGE_COMPUTATION) {
+        return PSA_ERROR_BAD_STATE;
+    }
+
     if (operation->id == 0) {
         return PSA_ERROR_BAD_STATE;
     }
@@ -7341,8 +7449,10 @@
 psa_status_t psa_pake_abort(
     psa_pake_operation_t *operation)
 {
-    /* Aborting a non-active operation is allowed */
-    if (operation->id == 0) {
+    /* If we are in collecting inputs stage clear inputs. */
+    if (operation->stage == PSA_PAKE_OPERATION_STAGE_COLLECT_INPUTS) {
+        mbedtls_free(operation->data.inputs.password);
+        memset(&operation->data.inputs, 0, sizeof(psa_crypto_driver_pake_inputs_t));
         return PSA_SUCCESS;
     }
 
diff --git a/library/psa_crypto_driver_wrappers.h b/library/psa_crypto_driver_wrappers.h
index 78f2f9a..abaabb5 100644
--- a/library/psa_crypto_driver_wrappers.h
+++ b/library/psa_crypto_driver_wrappers.h
@@ -417,27 +417,7 @@
  */
 psa_status_t psa_driver_wrapper_pake_setup(
     psa_pake_operation_t *operation,
-    const psa_pake_cipher_suite_t *cipher_suite);
-
-psa_status_t psa_driver_wrapper_pake_set_password_key(
-    const psa_key_attributes_t *attributes,
-    psa_pake_operation_t *operation,
-    uint8_t *key_buffer,
-    size_t key_size);
-
-psa_status_t psa_driver_wrapper_pake_set_user(
-    psa_pake_operation_t *operation,
-    const uint8_t *user_id,
-    size_t user_id_len);
-
-psa_status_t psa_driver_wrapper_pake_set_peer(
-    psa_pake_operation_t *operation,
-    const uint8_t *peer_id,
-    size_t peer_id_len);
-
-psa_status_t psa_driver_wrapper_pake_set_role(
-    psa_pake_operation_t *operation,
-    psa_pake_role_t role);
+    const psa_crypto_driver_pake_inputs_t *inputs);
 
 psa_status_t psa_driver_wrapper_pake_output(
     psa_pake_operation_t *operation,
diff --git a/library/psa_crypto_pake.c b/library/psa_crypto_pake.c
index 1e5dca4..3a710dc 100644
--- a/library/psa_crypto_pake.c
+++ b/library/psa_crypto_pake.c
@@ -192,36 +192,32 @@
 
 #if defined(MBEDTLS_PSA_BUILTIN_PAKE)
 psa_status_t mbedtls_psa_pake_setup(mbedtls_psa_pake_operation_t *operation,
-                                    const psa_pake_cipher_suite_t *cipher_suite)
+                                    const psa_crypto_driver_pake_inputs_t *inputs)
 {
     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;
-    }
+    uint8_t *password = inputs->password;
+    size_t password_len = inputs->password_len;
+    psa_pake_role_t role = inputs->role;
+    psa_pake_cipher_suite_t cipher_suite = inputs->cipher_suite;
 
-    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;
-    }
+    memset(operation, 0, sizeof(mbedtls_psa_pake_operation_t));
 
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
-    if (cipher_suite->algorithm == PSA_ALG_JPAKE) {
-        if (cipher_suite->type != PSA_PAKE_PRIMITIVE_TYPE_ECC ||
-            cipher_suite->family != PSA_ECC_FAMILY_SECP_R1 ||
-            cipher_suite->bits != 256 ||
-            cipher_suite->hash != PSA_ALG_SHA_256) {
+    if (cipher_suite.algorithm == PSA_ALG_JPAKE) {
+        if (cipher_suite.type != PSA_PAKE_PRIMITIVE_TYPE_ECC ||
+            cipher_suite.family != PSA_ECC_FAMILY_SECP_R1 ||
+            cipher_suite.bits != 256 ||
+            cipher_suite.hash != PSA_ALG_SHA_256) {
             status = PSA_ERROR_NOT_SUPPORTED;
             goto error;
         }
 
-        operation->alg = cipher_suite->algorithm;
+        if (role != PSA_PAKE_ROLE_CLIENT &&
+            role != PSA_PAKE_ROLE_SERVER) {
+            status = PSA_ERROR_NOT_SUPPORTED;
+            goto error;
+        }
 
         mbedtls_ecjpake_init(&operation->ctx.pake);
 
@@ -229,8 +225,10 @@
         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;
+        operation->password_len = password_len;
+        operation->password = password;
+        operation->role = role;
+        operation->alg = cipher_suite.algorithm;
 
         mbedtls_platform_zeroize(operation->buffer, MBEDTLS_PSA_PAKE_BUFFER_SIZE);
         operation->buffer_length = 0;
@@ -240,149 +238,16 @@
     } else
 #else
     (void) operation;
-    (void) cipher_suite;
+    (void) inputs;
 #endif
     { status = PSA_ERROR_NOT_SUPPORTED; }
 
 error:
+    mbedtls_free(password);
     mbedtls_psa_pake_abort(operation);
     return status;
 }
 
-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_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) {
-        status = PSA_ERROR_INVALID_ARGUMENT;
-        goto error;
-    }
-
-    if ((usage & PSA_KEY_USAGE_DERIVE) == 0) {
-        status = PSA_ERROR_NOT_PERMITTED;
-        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) {
-        status = PSA_ERROR_BAD_STATE;
-        goto error;
-    }
-
-    operation->password = mbedtls_calloc(1, password_len);
-    if (operation->password == NULL) {
-        return PSA_ERROR_INSUFFICIENT_MEMORY;
-    }
-
-    memcpy(operation->password, password, password_len);
-    operation->password_len = password_len;
-
-    return PSA_SUCCESS;
-
-error:
-    mbedtls_psa_pake_abort(operation);
-    return status;
-}
-
-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;
-    }
-
-    status = PSA_ERROR_NOT_SUPPORTED;
-
-error:
-    mbedtls_psa_pake_abort(operation);
-    return status;
-}
-
-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;
-    }
-
-    status = PSA_ERROR_NOT_SUPPORTED;
-
-error:
-    mbedtls_psa_pake_abort(operation);
-    return status;
-}
-
-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;
-    }
-
-#if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
-    if (operation->alg == PSA_ALG_JPAKE) {
-        if (role != PSA_PAKE_ROLE_CLIENT &&
-            role != PSA_PAKE_ROLE_SERVER) {
-            return PSA_ERROR_NOT_SUPPORTED;
-        }
-
-        operation->role = role;
-
-        return PSA_SUCCESS;
-    } else
-#else
-    (void) role;
-#endif
-
-    { status = PSA_ERROR_NOT_SUPPORTED; }
-
-error:
-    mbedtls_psa_pake_abort(operation);
-    return status;
-}
 
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
 static psa_status_t psa_pake_ecjpake_setup(mbedtls_psa_pake_operation_t *operation)
diff --git a/library/psa_crypto_pake.h b/library/psa_crypto_pake.h
index 9256f5a..4768cee 100644
--- a/library/psa_crypto_pake.h
+++ b/library/psa_crypto_pake.h
@@ -94,178 +94,8 @@
  *         results in this error code.
  */
 psa_status_t mbedtls_psa_pake_setup(mbedtls_psa_pake_operation_t *operation,
-                                    const psa_pake_cipher_suite_t *cipher_suite);
+                                    const psa_crypto_driver_pake_inputs_t *inputs);
 
-/** 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
- *                              psa_pake_input() has been called yet). It must
- *                              be on operation for which the password hasn't
- *                              been set yet (psa_pake_set_password_key()
- *                              hasn't been called yet).
- * \param password              Buffer holding the password
- * \param password_len          Password buffer size
- *
- * \retval #PSA_SUCCESS
- *         Success.
- * \retval #PSA_ERROR_INVALID_HANDLE
- *         \p password is not a valid key identifier.
- * \retval #PSA_ERROR_NOT_PERMITTED
- *         The key does not have the #PSA_KEY_USAGE_DERIVE flag, or it does not
- *         permit the \p operation's algorithm.
- * \retval #PSA_ERROR_INVALID_ARGUMENT
- *         The key type for \p password is not #PSA_KEY_TYPE_PASSWORD or
- *         #PSA_KEY_TYPE_PASSWORD_HASH, or \p password is not compatible with
- *         the \p operation's cipher suite.
- * \retval #PSA_ERROR_NOT_SUPPORTED
- *         The key type or key size of \p password is not supported with the
- *         \p operation's cipher suite.
- * \retval #PSA_ERROR_COMMUNICATION_FAILURE
- * \retval #PSA_ERROR_CORRUPTION_DETECTED
- * \retval #PSA_ERROR_STORAGE_FAILURE
- * \retval #PSA_ERROR_DATA_CORRUPT
- * \retval #PSA_ERROR_DATA_INVALID
- * \retval #PSA_ERROR_BAD_STATE
- *         The operation state is not valid (it must have been set up.), or
- *         the library has not been previously initialized by psa_crypto_init().
- *         It is implementation-dependent whether a failure to initialize
- *         results in this error code.
- */
-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);
-
-/** Set the user ID for a password-authenticated key exchange.
- *
- * Call this function to set the user ID. For PAKE algorithms that associate a
- * user identifier with each side of the session you need to call
- * psa_pake_set_peer() as well. For PAKE algorithms that associate a single
- * user identifier with the session, call psa_pake_set_user() only.
- *
- * Refer to the documentation of individual PAKE algorithm types (`PSA_ALG_XXX`
- * values of type ::psa_algorithm_t such that #PSA_ALG_IS_PAKE(\c alg) is true)
- * for more information.
- *
- * \param[in,out] operation     The operation object to set the user ID for. It
- *                              must have been set up by psa_pake_setup() and
- *                              not yet in use (neither psa_pake_output() nor
- *                              psa_pake_input() has been called yet). It must
- *                              be on operation for which the user ID hasn't
- *                              been set (psa_pake_set_user() hasn't been
- *                              called yet).
- * \param[in] user_id           The user ID to authenticate with.
- * \param user_id_len           Size of the \p user_id buffer in bytes.
- *
- * \retval #PSA_SUCCESS
- *         Success.
- * \retval #PSA_ERROR_INVALID_ARGUMENT
- *         \p user_id is not valid for the \p operation's algorithm and cipher
- *         suite.
- * \retval #PSA_ERROR_NOT_SUPPORTED
- *         The value of \p user_id is not supported by the implementation.
- * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
- * \retval #PSA_ERROR_COMMUNICATION_FAILURE
- * \retval #PSA_ERROR_CORRUPTION_DETECTED
- * \retval #PSA_ERROR_BAD_STATE
- *         The operation state is not valid, or
- *         the library has not been previously initialized by psa_crypto_init().
- *         It is implementation-dependent whether a failure to initialize
- *         results in this error code.
- */
-psa_status_t mbedtls_psa_pake_set_user(mbedtls_psa_pake_operation_t *operation,
-                                       const uint8_t *user_id,
-                                       size_t user_id_len);
-
-/** Set the peer ID for a password-authenticated key exchange.
- *
- * Call this function in addition to psa_pake_set_user() for PAKE algorithms
- * that associate a user identifier with each side of the session. For PAKE
- * algorithms that associate a single user identifier with the session, call
- * psa_pake_set_user() only.
- *
- * Refer to the documentation of individual PAKE algorithm types (`PSA_ALG_XXX`
- * values of type ::psa_algorithm_t such that #PSA_ALG_IS_PAKE(\c alg) is true)
- * for more information.
- *
- * \param[in,out] operation     The operation object to set the peer ID for. It
- *                              must have been set up by psa_pake_setup() and
- *                              not yet in use (neither psa_pake_output() nor
- *                              psa_pake_input() has been called yet). It must
- *                              be on operation for which the peer ID hasn't
- *                              been set (psa_pake_set_peer() hasn't been
- *                              called yet).
- * \param[in] peer_id           The peer's ID to authenticate.
- * \param peer_id_len           Size of the \p peer_id buffer in bytes.
- *
- * \retval #PSA_SUCCESS
- *         Success.
- * \retval #PSA_ERROR_INVALID_ARGUMENT
- *         \p user_id is not valid for the \p operation's algorithm and cipher
- *         suite.
- * \retval #PSA_ERROR_NOT_SUPPORTED
- *         The algorithm doesn't associate a second identity with the session.
- * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
- * \retval #PSA_ERROR_COMMUNICATION_FAILURE
- * \retval #PSA_ERROR_CORRUPTION_DETECTED
- * \retval #PSA_ERROR_BAD_STATE
- *         Calling psa_pake_set_peer() is invalid with the \p operation's
- *         algorithm, the operation state is not valid, or the library has not
- *         been previously initialized by psa_crypto_init().
- *         It is implementation-dependent whether a failure to initialize
- *         results in this error code.
- */
-psa_status_t mbedtls_psa_pake_set_peer(mbedtls_psa_pake_operation_t *operation,
-                                       const uint8_t *peer_id,
-                                       size_t peer_id_len);
-
-/** Set the application role for a password-authenticated key exchange.
- *
- * Not all PAKE algorithms need to differentiate the communicating entities.
- * It is optional to call this function for PAKEs that don't require a role
- * to be specified. For such PAKEs the application role parameter is ignored,
- * or #PSA_PAKE_ROLE_NONE can be passed as \c role.
- *
- * Refer to the documentation of individual PAKE algorithm types (`PSA_ALG_XXX`
- * values of type ::psa_algorithm_t such that #PSA_ALG_IS_PAKE(\c alg) is true)
- * for more information.
- *
- * \param[in,out] operation     The operation object to specify the
- *                              application's role for. It must have been set up
- *                              by psa_pake_setup() and not yet in use (neither
- *                              psa_pake_output() nor psa_pake_input() has been
- *                              called yet). It must be on operation for which
- *                              the application's role hasn't been specified
- *                              (psa_pake_set_role() hasn't been called yet).
- * \param role                  A value of type ::psa_pake_role_t indicating the
- *                              application's role in the PAKE the algorithm
- *                              that is being set up. For more information see
- *                              the documentation of \c PSA_PAKE_ROLE_XXX
- *                              constants.
- *
- * \retval #PSA_SUCCESS
- *         Success.
- * \retval #PSA_ERROR_INVALID_ARGUMENT
- *         The \p role is not a valid PAKE role in the \p operation’s algorithm.
- * \retval #PSA_ERROR_NOT_SUPPORTED
- *         The \p role for this algorithm is not supported or is not valid.
- * \retval #PSA_ERROR_COMMUNICATION_FAILURE
- * \retval #PSA_ERROR_CORRUPTION_DETECTED
- * \retval #PSA_ERROR_BAD_STATE
- *         The operation state is not valid, or
- *         the library has not been previously initialized by psa_crypto_init().
- *         It is implementation-dependent whether a failure to initialize
- *         results in this error code.
- */
-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.
  *