diff --git a/include/psa/crypto_extra.h b/include/psa/crypto_extra.h
index 216039c..66e5dbc 100644
--- a/include/psa/crypto_extra.h
+++ b/include/psa/crypto_extra.h
@@ -159,7 +159,7 @@
  *
  * \param[in,out] generator       The generator object to set up. It must have
  *                                been initialized as per the documentation for
- *                                #psa_crypto_generator_t and not yet in use.
+ *                                #psa_key_derivation_operation_t and not yet in use.
  * \param handle                  Handle to the secret key.
  * \param alg                     The key derivation algorithm to compute
  *                                (\c PSA_ALG_XXX value such that
@@ -190,7 +190,7 @@
  *         It is implementation-dependent whether a failure to initialize
  *         results in this error code.
  */
-psa_status_t psa_key_derivation(psa_crypto_generator_t *generator,
+psa_status_t psa_key_derivation(psa_key_derivation_operation_t *generator,
                                 psa_key_handle_t handle,
                                 psa_algorithm_t alg,
                                 const uint8_t *salt,
@@ -433,7 +433,7 @@
 psa_status_t psa_generate_derived_key_to_handle(psa_key_handle_t handle,
                                       psa_key_type_t type,
                                       size_t bits,
-                                      psa_crypto_generator_t *generator);
+                                      psa_key_derivation_operation_t *generator);
 
 psa_status_t psa_generate_random_key_to_handle(psa_key_handle_t handle,
                               psa_key_type_t type,