diff --git a/include/psa/crypto_extra.h b/include/psa/crypto_extra.h
index da74bed..4fa273d 100644
--- a/include/psa/crypto_extra.h
+++ b/include/psa/crypto_extra.h
@@ -429,6 +429,9 @@
  */
 #define PSA_DH_FAMILY_CUSTOM             ((psa_dh_family_t) 0x7e)
 
+/** EC-JPAKE operation stages. */
+#define PSA_PAKE_OPERATION_STAGE_COLLECT_INPUTS 0
+#define PSA_PAKE_OPERATION_STAGE_COMPUTATION 1
 
 /**
  * \brief Set domain parameters for a key.
@@ -1286,6 +1289,9 @@
  * Implementation details can change in future versions without notice. */
 typedef struct psa_pake_operation_s psa_pake_operation_t;
 
+/** The type of input values for PAKE operations. */
+typedef struct psa_crypto_driver_pake_inputs_s psa_crypto_driver_pake_inputs_t;
+
 /** Return an initial value for a PAKE operation object.
  */
 static psa_pake_operation_t psa_pake_operation_init(void);
@@ -1826,7 +1832,7 @@
 /** Returns a suitable initializer for a PAKE operation object of type
  * psa_pake_operation_t.
  */
-#define PSA_PAKE_OPERATION_INIT { 0, { .dummy = 0 } }
+#define PSA_PAKE_OPERATION_INIT { 0, PSA_PAKE_OPERATION_STAGE_COLLECT_INPUTS, { 0 } }
 
 struct psa_pake_cipher_suite_s {
     psa_algorithm_t algorithm;
@@ -1897,6 +1903,15 @@
     }
 }
 
+struct psa_crypto_driver_pake_inputs_s {
+    psa_algorithm_t MBEDTLS_PRIVATE(alg);
+    uint8_t *MBEDTLS_PRIVATE(password);
+    size_t MBEDTLS_PRIVATE(password_len);
+    psa_pake_role_t MBEDTLS_PRIVATE(role);
+    psa_key_lifetime_t MBEDTLS_PRIVATE(key_lifetime);
+    psa_pake_cipher_suite_t MBEDTLS_PRIVATE(cipher_suite);
+};
+
 struct psa_pake_operation_s {
     /** Unique ID indicating which driver got assigned to do the
      * operation. Since driver contexts are driver-specific, swapping
@@ -1905,7 +1920,15 @@
      * ID value zero means the context is not valid or not assigned to
      * any driver (i.e. none of the driver contexts are active). */
     unsigned int MBEDTLS_PRIVATE(id);
-    psa_driver_pake_context_t MBEDTLS_PRIVATE(ctx);
+    /* Based on stage (collecting inputs/computation) we select active structure of data union.
+     * While switching stage (when driver setup is called) collected inputs
+       are copied to the corresponding operation context. */
+    uint8_t MBEDTLS_PRIVATE(stage);
+    union {
+        unsigned dummy;
+        psa_crypto_driver_pake_inputs_t MBEDTLS_PRIVATE(inputs);
+        psa_driver_pake_context_t MBEDTLS_PRIVATE(ctx);
+    } MBEDTLS_PRIVATE(data);
 };
 
 static inline struct psa_pake_cipher_suite_s psa_pake_cipher_suite_init(void)
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.
  *
diff --git a/scripts/data_files/driver_templates/psa_crypto_driver_wrappers.c.jinja b/scripts/data_files/driver_templates/psa_crypto_driver_wrappers.c.jinja
index 9a7b645..21a3b5f 100644
--- a/scripts/data_files/driver_templates/psa_crypto_driver_wrappers.c.jinja
+++ b/scripts/data_files/driver_templates/psa_crypto_driver_wrappers.c.jinja
@@ -2811,64 +2811,12 @@
 
 psa_status_t psa_driver_wrapper_pake_setup(
     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;
 
-    /* Try setup on accelerators first */
-#if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
-#if defined(PSA_CRYPTO_DRIVER_TEST)
-    status = mbedtls_test_transparent_pake_setup(
-                        &operation->ctx.transparent_test_driver_ctx,
-                        (const psa_pake_cipher_suite_t*) cipher_suite );
-    if( status == PSA_SUCCESS )
-        operation->id = MBEDTLS_TEST_TRANSPARENT_DRIVER_ID;
-
-    if( status != PSA_ERROR_NOT_SUPPORTED )
-        return( status );
-#endif /* PSA_CRYPTO_DRIVER_TEST */
-#endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
-
-    /* If software fallback is compiled in, try fallback */
-#if defined(MBEDTLS_PSA_BUILTIN_PAKE)
-    status = mbedtls_psa_pake_setup( &operation->ctx.mbedtls_ctx, cipher_suite );
-    if( status == PSA_SUCCESS )
-        operation->id = PSA_CRYPTO_MBED_TLS_DRIVER_ID;
-
-    if( status != PSA_ERROR_NOT_SUPPORTED )
-        return( status );
-#endif /* MBEDTLS_PSA_BUILTIN_PAKE */
-
-    /* Add cases for opaque driver here */
-#if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
-#if defined(PSA_CRYPTO_DRIVER_TEST)
-    status = mbedtls_test_opaque_pake_setup(
-                        &operation->ctx.opaque_test_driver_ctx,
-                        (const psa_pake_cipher_suite_t*) cipher_suite );
-    if( status == PSA_SUCCESS )
-        operation->id = MBEDTLS_TEST_OPAQUE_DRIVER_ID;
-
-    if( status != PSA_ERROR_NOT_SUPPORTED )
-        return( status );
-#endif /* PSA_CRYPTO_DRIVER_TEST */
-#endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
-
-    /* Nothing left to try if we fall through here */
-    (void) status;
-    (void) operation;
-    (void) cipher_suite;
-    return( PSA_ERROR_NOT_SUPPORTED );
-}
-
-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 status = PSA_ERROR_CORRUPTION_DETECTED;
     psa_key_location_t location =
-        PSA_KEY_LIFETIME_GET_LOCATION( attributes->core.lifetime );
+            PSA_KEY_LIFETIME_GET_LOCATION( inputs->key_lifetime );
 
     switch( location )
     {
@@ -2877,135 +2825,44 @@
              * cycle through all known transparent accelerators */
 #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
 #if defined(PSA_CRYPTO_DRIVER_TEST)
-            status = mbedtls_test_transparent_set_password_key(
-                        attributes,
-                        &operation->ctx.transparent_test_driver_ctx,
-                        key_buffer, key_size );
+            status = mbedtls_test_transparent_pake_setup(
+                        &operation->data.ctx.transparent_test_driver_ctx,
+                        inputs );
+            if( status == PSA_SUCCESS )
+                operation->id = MBEDTLS_TEST_TRANSPARENT_DRIVER_ID;
             /* Declared with fallback == true */
             if( status != PSA_ERROR_NOT_SUPPORTED )
                 return( status );
 #endif /* PSA_CRYPTO_DRIVER_TEST */
 #endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
 #if defined(MBEDTLS_PSA_BUILTIN_PAKE)
-            return( mbedtls_psa_pake_set_password_key(
-                        attributes, &operation->ctx.mbedtls_ctx,
-                        key_buffer, key_size ) );
+            status = mbedtls_psa_pake_setup( &operation->data.ctx.mbedtls_ctx,
+                        inputs );
+            if( status == PSA_SUCCESS )
+                operation->id = PSA_CRYPTO_MBED_TLS_DRIVER_ID;
+            return status;
 #endif
             return( PSA_ERROR_NOT_SUPPORTED );
         /* Add cases for opaque driver here */
 #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
 #if defined(PSA_CRYPTO_DRIVER_TEST)
         case PSA_CRYPTO_TEST_DRIVER_LOCATION:
-            return( mbedtls_test_opaque_set_password_key(
-                        attributes,
-                        &operation->ctx.opaque_test_driver_ctx,
-                        key_buffer, key_size ) );
+            status = mbedtls_test_opaque_pake_setup(
+                        &operation->data.ctx.opaque_test_driver_ctx,
+                        inputs );
+             if( status == PSA_SUCCESS )
+                operation->id = MBEDTLS_TEST_OPAQUE_DRIVER_ID;
+            return status;
 #endif /* PSA_CRYPTO_DRIVER_TEST */
 #endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
-
         default:
             /* Key is declared with a lifetime not known to us */
             (void)operation;
+            (void)inputs;
             (void)status;
-            (void)key_buffer;
-            (void)key_size;
             return( PSA_ERROR_INVALID_ARGUMENT );
     }
 }
-
-psa_status_t psa_driver_wrapper_pake_set_user(
-    psa_pake_operation_t *operation,
-    const uint8_t *user_id,
-    size_t user_id_len )
-{
-    switch( operation->id )
-    {
-#if defined(MBEDTLS_PSA_BUILTIN_PAKE)
-        case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
-            return( mbedtls_psa_pake_set_user( &operation->ctx.mbedtls_ctx,
-                                               user_id, user_id_len ) );
-#endif /* MBEDTLS_PSA_BUILTIN_PAKE */
-
-#if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
-#if defined(PSA_CRYPTO_DRIVER_TEST)
-        case MBEDTLS_TEST_TRANSPARENT_DRIVER_ID:
-            return( mbedtls_test_transparent_pake_set_user(
-                        &operation->ctx.transparent_test_driver_ctx,
-                        user_id, user_id_len ) );
-        case MBEDTLS_TEST_OPAQUE_DRIVER_ID:
-            return( mbedtls_test_opaque_pake_set_user(
-                        &operation->ctx.opaque_test_driver_ctx,
-                        user_id, user_id_len ) );
-#endif /* PSA_CRYPTO_DRIVER_TEST */
-#endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
-        default:
-            (void) user_id;
-            (void) user_id_len;
-            return( PSA_ERROR_INVALID_ARGUMENT );
-    }
-}
-
-psa_status_t psa_driver_wrapper_pake_set_peer(
-    psa_pake_operation_t *operation,
-    const uint8_t *peer_id,
-    size_t peer_id_len )
-{
-    switch( operation->id )
-    {
-#if defined(MBEDTLS_PSA_BUILTIN_PAKE)
-        case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
-            return( mbedtls_psa_pake_set_peer( &operation->ctx.mbedtls_ctx,
-                                               peer_id, peer_id_len ) );
-#endif /* MBEDTLS_PSA_BUILTIN_PAKE */
-
-#if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
-#if defined(PSA_CRYPTO_DRIVER_TEST)
-        case MBEDTLS_TEST_TRANSPARENT_DRIVER_ID:
-            return( mbedtls_test_transparent_pake_set_peer(
-                        &operation->ctx.transparent_test_driver_ctx,
-                        peer_id, peer_id_len ) );
-        case MBEDTLS_TEST_OPAQUE_DRIVER_ID:
-            return( mbedtls_test_opaque_pake_set_peer(
-                        &operation->ctx.opaque_test_driver_ctx,
-                        peer_id, peer_id_len ) );
-#endif /* PSA_CRYPTO_DRIVER_TEST */
-#endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
-        default:
-            (void) peer_id;
-            (void) peer_id_len;
-            return( PSA_ERROR_INVALID_ARGUMENT );
-    }
-}
-
-psa_status_t psa_driver_wrapper_pake_set_role(
-    psa_pake_operation_t *operation,
-    psa_pake_role_t role )
-{
-    switch( operation->id )
-    {
-#if defined(MBEDTLS_PSA_BUILTIN_PAKE)
-        case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
-            return( mbedtls_psa_pake_set_role( &operation->ctx.mbedtls_ctx, role ) );
-#endif /* MBEDTLS_PSA_BUILTIN_PAKE */
-
-#if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
-#if defined(PSA_CRYPTO_DRIVER_TEST)
-        case MBEDTLS_TEST_TRANSPARENT_DRIVER_ID:
-            return( mbedtls_test_transparent_pake_set_role(
-                        &operation->ctx.transparent_test_driver_ctx,
-                        role ) );
-        case MBEDTLS_TEST_OPAQUE_DRIVER_ID:
-            return( mbedtls_test_opaque_pake_set_role(
-                        &operation->ctx.opaque_test_driver_ctx,
-                        role ) );
-#endif /* PSA_CRYPTO_DRIVER_TEST */
-#endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
-        default:
-            (void) role;
-            return( PSA_ERROR_INVALID_ARGUMENT );
-    }
-}
-
 psa_status_t psa_driver_wrapper_pake_output(
     psa_pake_operation_t *operation,
     psa_pake_step_t step,
@@ -3017,7 +2874,7 @@
     {
 #if defined(MBEDTLS_PSA_BUILTIN_PAKE)
         case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
-            return( mbedtls_psa_pake_output( &operation->ctx.mbedtls_ctx, step, output,
+            return( mbedtls_psa_pake_output( &operation->data.ctx.mbedtls_ctx, step, output,
                                              output_size, output_length ) );
 #endif /* MBEDTLS_PSA_BUILTIN_PAKE */
 
@@ -3025,11 +2882,11 @@
 #if defined(PSA_CRYPTO_DRIVER_TEST)
         case MBEDTLS_TEST_TRANSPARENT_DRIVER_ID:
             return( mbedtls_test_transparent_pake_output(
-                        &operation->ctx.transparent_test_driver_ctx,
+                        &operation->data.ctx.transparent_test_driver_ctx,
                         step, output, output_size, output_length ) );
         case MBEDTLS_TEST_OPAQUE_DRIVER_ID:
             return( mbedtls_test_opaque_pake_output(
-                        &operation->ctx.opaque_test_driver_ctx,
+                        &operation->data.ctx.opaque_test_driver_ctx,
                         step, output, output_size, output_length ) );
 #endif /* PSA_CRYPTO_DRIVER_TEST */
 #endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
@@ -3052,7 +2909,7 @@
     {
 #if defined(MBEDTLS_PSA_BUILTIN_PAKE)
         case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
-            return( mbedtls_psa_pake_input( &operation->ctx.mbedtls_ctx,
+            return( mbedtls_psa_pake_input( &operation->data.ctx.mbedtls_ctx,
                                             step, input, input_length ) );
 #endif /* MBEDTLS_PSA_BUILTIN_PAKE */
 
@@ -3060,11 +2917,11 @@
 #if defined(PSA_CRYPTO_DRIVER_TEST)
         case MBEDTLS_TEST_TRANSPARENT_DRIVER_ID:
             return( mbedtls_test_transparent_pake_input(
-                        &operation->ctx.transparent_test_driver_ctx,
+                        &operation->data.ctx.transparent_test_driver_ctx,
                         step, input, input_length ) );
         case MBEDTLS_TEST_OPAQUE_DRIVER_ID:
             return( mbedtls_test_opaque_pake_input(
-                        &operation->ctx.opaque_test_driver_ctx,
+                        &operation->data.ctx.opaque_test_driver_ctx,
                         step, input, input_length ) );
 #endif /* PSA_CRYPTO_DRIVER_TEST */
 #endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
@@ -3084,18 +2941,18 @@
     {
 #if defined(MBEDTLS_PSA_BUILTIN_PAKE)
         case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
-            return( mbedtls_psa_pake_get_implicit_key( &operation->ctx.mbedtls_ctx, output, output_size ) );
+            return( mbedtls_psa_pake_get_implicit_key( &operation->data.ctx.mbedtls_ctx, output, output_size ) );
 #endif /* MBEDTLS_PSA_BUILTIN_PAKE */
 
 #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
 #if defined(PSA_CRYPTO_DRIVER_TEST)
         case MBEDTLS_TEST_TRANSPARENT_DRIVER_ID:
             return( mbedtls_test_transparent_pake_get_implicit_key(
-                        &operation->ctx.transparent_test_driver_ctx,
+                        &operation->data.ctx.transparent_test_driver_ctx,
                         output, output_size ) );
         case MBEDTLS_TEST_OPAQUE_DRIVER_ID:
             return( mbedtls_test_opaque_pake_get_implicit_key(
-                        &operation->ctx.opaque_test_driver_ctx,
+                        &operation->data.ctx.opaque_test_driver_ctx,
                         output, output_size ) );
 #endif /* PSA_CRYPTO_DRIVER_TEST */
 #endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
@@ -3113,17 +2970,17 @@
     {
 #if defined(MBEDTLS_PSA_BUILTIN_PAKE)
         case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
-            return( mbedtls_psa_pake_abort( &operation->ctx.mbedtls_ctx ) );
+            return( mbedtls_psa_pake_abort( &operation->data.ctx.mbedtls_ctx ) );
 #endif /* MBEDTLS_PSA_BUILTIN_PAKE */
 
 #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
 #if defined(PSA_CRYPTO_DRIVER_TEST)
         case MBEDTLS_TEST_TRANSPARENT_DRIVER_ID:
             return( mbedtls_test_transparent_pake_abort(
-                        &operation->ctx.transparent_test_driver_ctx ) );
+                        &operation->data.ctx.transparent_test_driver_ctx ) );
         case MBEDTLS_TEST_OPAQUE_DRIVER_ID:
             return( mbedtls_test_opaque_pake_abort(
-                        &operation->ctx.opaque_test_driver_ctx ) );
+                        &operation->data.ctx.opaque_test_driver_ctx ) );
 #endif /* PSA_CRYPTO_DRIVER_TEST */
 #endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
         default:
diff --git a/tests/include/test/drivers/pake.h b/tests/include/test/drivers/pake.h
index 5ee401b..b1d3d44 100644
--- a/tests/include/test/drivers/pake.h
+++ b/tests/include/test/drivers/pake.h
@@ -50,27 +50,7 @@
 
 psa_status_t mbedtls_test_transparent_pake_setup(
     mbedtls_transparent_test_driver_pake_operation_t *operation,
-    const psa_pake_cipher_suite_t *cipher_suite);
-
-psa_status_t mbedtls_test_transparent_set_password_key(
-    const psa_key_attributes_t *attributes,
-    mbedtls_transparent_test_driver_pake_operation_t *operation,
-    uint8_t *key_buffer,
-    size_t key_size);
-
-psa_status_t mbedtls_test_transparent_pake_set_user(
-    mbedtls_transparent_test_driver_pake_operation_t *operation,
-    const uint8_t *user_id,
-    size_t user_id_len);
-
-psa_status_t mbedtls_test_transparent_pake_set_peer(
-    mbedtls_transparent_test_driver_pake_operation_t *operation,
-    const uint8_t *peer_id,
-    size_t peer_id_len);
-
-psa_status_t mbedtls_test_transparent_pake_set_role(
-    mbedtls_transparent_test_driver_pake_operation_t *operation,
-    psa_pake_role_t role);
+    const psa_crypto_driver_pake_inputs_t *inputs);
 
 psa_status_t mbedtls_test_transparent_pake_output(
     mbedtls_transparent_test_driver_pake_operation_t *operation,
@@ -94,7 +74,7 @@
 
 psa_status_t mbedtls_test_opaque_pake_setup(
     mbedtls_opaque_test_driver_pake_operation_t *operation,
-    const psa_pake_cipher_suite_t *cipher_suite);
+    const psa_crypto_driver_pake_inputs_t *inputs);
 
 psa_status_t mbedtls_test_opaque_set_password_key(
     const psa_key_attributes_t *attributes,
diff --git a/tests/src/drivers/test_driver_pake.c b/tests/src/drivers/test_driver_pake.c
index 3495705..06168a1 100644
--- a/tests/src/drivers/test_driver_pake.c
+++ b/tests/src/drivers/test_driver_pake.c
@@ -35,7 +35,7 @@
 
 psa_status_t mbedtls_test_transparent_pake_setup(
     mbedtls_transparent_test_driver_pake_operation_t *operation,
-    const psa_pake_cipher_suite_t *cipher_suite)
+    const psa_crypto_driver_pake_inputs_t *inputs)
 {
     mbedtls_test_driver_pake_hooks.hits++;
 
@@ -47,139 +47,14 @@
         defined(LIBTESTDRIVER1_MBEDTLS_PSA_BUILTIN_PAKE)
         mbedtls_test_driver_pake_hooks.driver_status =
             libtestdriver1_mbedtls_psa_pake_setup(
-                operation, (const libtestdriver1_psa_pake_cipher_suite_t *) cipher_suite);
+                operation, (const libtestdriver1_psa_crypto_driver_pake_inputs_t *) inputs);
 #elif defined(MBEDTLS_PSA_BUILTIN_PAKE)
         mbedtls_test_driver_pake_hooks.driver_status =
             mbedtls_psa_pake_setup(
-                operation, cipher_suite);
+                operation, inputs);
 #else
         (void) operation;
-        (void) cipher_suite;
-        mbedtls_test_driver_pake_hooks.driver_status = PSA_ERROR_NOT_SUPPORTED;
-#endif
-    }
-
-    return mbedtls_test_driver_pake_hooks.driver_status;
-}
-
-psa_status_t mbedtls_test_transparent_set_password_key(
-    const psa_key_attributes_t *attributes,
-    mbedtls_transparent_test_driver_pake_operation_t *operation,
-    uint8_t *key_buffer,
-    size_t key_size)
-{
-    mbedtls_test_driver_pake_hooks.hits++;
-
-    if (mbedtls_test_driver_pake_hooks.forced_status != PSA_SUCCESS) {
-        mbedtls_test_driver_pake_hooks.driver_status =
-            mbedtls_test_driver_pake_hooks.forced_status;
-    } else {
-#if defined(MBEDTLS_TEST_LIBTESTDRIVER1) && \
-        defined(LIBTESTDRIVER1_MBEDTLS_PSA_BUILTIN_PAKE)
-        mbedtls_test_driver_pake_hooks.driver_status =
-            libtestdriver1_mbedtls_psa_pake_set_password_key(
-                (const libtestdriver1_psa_key_attributes_t *) attributes,
-                operation, key_buffer, key_size);
-#elif defined(MBEDTLS_PSA_BUILTIN_PAKE)
-        mbedtls_test_driver_pake_hooks.driver_status =
-            mbedtls_psa_pake_set_password_key(
-                attributes, operation, key_buffer, key_size);
-#else
-        (void) operation;
-        (void) key_buffer,
-        (void) key_size;
-        mbedtls_test_driver_pake_hooks.driver_status = PSA_ERROR_NOT_SUPPORTED;
-#endif
-    }
-
-    return mbedtls_test_driver_pake_hooks.driver_status;
-}
-
-psa_status_t mbedtls_test_transparent_pake_set_user(
-    mbedtls_transparent_test_driver_pake_operation_t *operation,
-    const uint8_t *user_id,
-    size_t user_id_len)
-{
-    mbedtls_test_driver_pake_hooks.hits++;
-
-    if (mbedtls_test_driver_pake_hooks.forced_status != PSA_SUCCESS) {
-        mbedtls_test_driver_pake_hooks.driver_status =
-            mbedtls_test_driver_pake_hooks.forced_status;
-    } else {
-#if defined(MBEDTLS_TEST_LIBTESTDRIVER1) && \
-        defined(LIBTESTDRIVER1_MBEDTLS_PSA_BUILTIN_PAKE)
-        mbedtls_test_driver_pake_hooks.driver_status =
-            libtestdriver1_mbedtls_psa_pake_set_user(
-                operation, user_id, user_id_len);
-#elif defined(MBEDTLS_PSA_BUILTIN_PAKE)
-        mbedtls_test_driver_pake_hooks.driver_status =
-            mbedtls_psa_pake_set_user(
-                operation, user_id, user_id_len);
-#else
-        (void) operation;
-        (void) user_id;
-        (void) user_id_len;
-        mbedtls_test_driver_pake_hooks.driver_status = PSA_ERROR_NOT_SUPPORTED;
-#endif
-    }
-
-    return mbedtls_test_driver_pake_hooks.driver_status;
-}
-
-
-psa_status_t mbedtls_test_transparent_pake_set_peer(
-    mbedtls_transparent_test_driver_pake_operation_t *operation,
-    const uint8_t *peer_id,
-    size_t peer_id_len)
-{
-    mbedtls_test_driver_pake_hooks.hits++;
-
-    if (mbedtls_test_driver_pake_hooks.forced_status != PSA_SUCCESS) {
-        mbedtls_test_driver_pake_hooks.driver_status =
-            mbedtls_test_driver_pake_hooks.forced_status;
-    } else {
-#if defined(MBEDTLS_TEST_LIBTESTDRIVER1) && \
-        defined(LIBTESTDRIVER1_MBEDTLS_PSA_BUILTIN_PAKE)
-        mbedtls_test_driver_pake_hooks.driver_status =
-            libtestdriver1_mbedtls_psa_pake_set_peer(
-                operation, peer_id, peer_id_len);
-#elif defined(MBEDTLS_PSA_BUILTIN_PAKE)
-        mbedtls_test_driver_pake_hooks.driver_status =
-            mbedtls_psa_pake_set_peer(
-                operation, peer_id, peer_id_len);
-#else
-        (void) operation;
-        (void) peer_id;
-        (void) peer_id_len;
-        mbedtls_test_driver_pake_hooks.driver_status = PSA_ERROR_NOT_SUPPORTED;
-#endif
-    }
-
-    return mbedtls_test_driver_pake_hooks.driver_status;
-}
-
-psa_status_t mbedtls_test_transparent_pake_set_role(
-    mbedtls_transparent_test_driver_pake_operation_t *operation,
-    psa_pake_role_t role)
-{
-    mbedtls_test_driver_pake_hooks.hits++;
-
-    if (mbedtls_test_driver_pake_hooks.forced_status != PSA_SUCCESS) {
-        mbedtls_test_driver_pake_hooks.driver_status =
-            mbedtls_test_driver_pake_hooks.forced_status;
-    } else {
-#if defined(MBEDTLS_TEST_LIBTESTDRIVER1) && \
-        defined(LIBTESTDRIVER1_MBEDTLS_PSA_BUILTIN_PAKE)
-        mbedtls_test_driver_pake_hooks.driver_status =
-            libtestdriver1_mbedtls_psa_pake_set_role(
-                operation, role);
-#elif defined(MBEDTLS_PSA_BUILTIN_PAKE)
-        mbedtls_test_driver_pake_hooks.driver_status =
-            mbedtls_psa_pake_set_role(
-                operation, role);
-#else
-        (void) operation;
-        (void) role;
+        (void) inputs;
         mbedtls_test_driver_pake_hooks.driver_status = PSA_ERROR_NOT_SUPPORTED;
 #endif
     }
@@ -329,10 +204,10 @@
  */
 psa_status_t mbedtls_test_opaque_pake_setup(
     mbedtls_opaque_test_driver_pake_operation_t *operation,
-    const psa_pake_cipher_suite_t *cipher_suite)
+    const psa_crypto_driver_pake_inputs_t *inputs)
 {
     (void) operation;
-    (void) cipher_suite;
+    (void) inputs;
     return PSA_ERROR_NOT_SUPPORTED;
 }