diff --git a/library/psa_crypto_pake.h b/library/psa_crypto_pake.h
index 485c93a..dc6ad7b 100644
--- a/library/psa_crypto_pake.h
+++ b/library/psa_crypto_pake.h
@@ -58,7 +58,6 @@
  * \param[in,out] operation    Active PAKE operation.
  * \param step                 The step of the algorithm for which the output is
  *                             requested.
- * \param computation_stage    The structure that holds PAKE computation stage.
  * \param[out] output          Buffer where the output is to be written in the
  *                             format appropriate for this \p step. Refer to
  *                             the documentation of the individual
@@ -97,8 +96,7 @@
  *         results in this error code.
  */
 psa_status_t mbedtls_psa_pake_output(mbedtls_psa_pake_operation_t *operation,
-                                     psa_pake_step_t step,
-                                     const psa_pake_computation_stage_t *computation_stage,
+                                     psa_pake_driver_step_t step,
                                      uint8_t *output,
                                      size_t output_size,
                                      size_t *output_length);
@@ -112,7 +110,6 @@
  *
  * \param[in,out] operation    Active PAKE operation.
  * \param step                 The step for which the input is provided.
- * \param computation_stage    The structure that holds PAKE computation stage.
  * \param[in] input            Buffer containing the input in the format
  *                             appropriate for this \p step. Refer to the
  *                             documentation of the individual
@@ -146,8 +143,7 @@
  *         results in this error code.
  */
 psa_status_t mbedtls_psa_pake_input(mbedtls_psa_pake_operation_t *operation,
-                                    psa_pake_step_t step,
-                                    const psa_pake_computation_stage_t *computation_stage,
+                                    psa_pake_driver_step_t step,
                                     const uint8_t *input,
                                     size_t input_length);