diff --git a/library/psa_crypto_pake.h b/library/psa_crypto_pake.h
index dc6ad7b..3658556 100644
--- a/library/psa_crypto_pake.h
+++ b/library/psa_crypto_pake.h
@@ -96,7 +96,7 @@
  *         results in this error code.
  */
 psa_status_t mbedtls_psa_pake_output(mbedtls_psa_pake_operation_t *operation,
-                                     psa_pake_driver_step_t step,
+                                     psa_crypto_driver_pake_step_t step,
                                      uint8_t *output,
                                      size_t output_size,
                                      size_t *output_length);
@@ -143,7 +143,7 @@
  *         results in this error code.
  */
 psa_status_t mbedtls_psa_pake_input(mbedtls_psa_pake_operation_t *operation,
-                                    psa_pake_driver_step_t step,
+                                    psa_crypto_driver_pake_step_t step,
                                     const uint8_t *input,
                                     size_t input_length);