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.
  *