diff --git a/include/psa/crypto.h b/include/psa/crypto.h
index 2217b95ed..0e842e5cb 100644
--- a/include/psa/crypto.h
+++ b/include/psa/crypto.h
@@ -2231,59 +2231,6 @@ psa_status_t psa_key_derivation_input_key(psa_crypto_generator_t *generator,
                                           psa_key_derivation_step_t step,
                                           psa_key_handle_t handle);
 
-/** Set up a key derivation operation.
- *
- * A key derivation algorithm takes three inputs: a secret input \p key and
- * two non-secret inputs \p label and p salt.
- * The result of this function is a byte generator which can
- * be used to produce keys and other cryptographic material.
- *
- * The role of \p label and \p salt is as follows:
- * - For HKDF (#PSA_ALG_HKDF), \p salt is the salt used in the "extract" step
- *   and \p label is the info string used in the "expand" step.
- *
- * \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.
- * \param handle                  Handle to the secret key.
- * \param alg                     The key derivation algorithm to compute
- *                                (\c PSA_ALG_XXX value such that
- *                                #PSA_ALG_IS_KEY_DERIVATION(\p alg) is true).
- * \param[in] salt                Salt to use.
- * \param salt_length             Size of the \p salt buffer in bytes.
- * \param[in] label               Label to use.
- * \param label_length            Size of the \p label buffer in bytes.
- * \param capacity                The maximum number of bytes that the
- *                                generator will be able to provide.
- *
- * \retval #PSA_SUCCESS
- *         Success.
- * \retval #PSA_ERROR_INVALID_HANDLE
- * \retval #PSA_ERROR_EMPTY_SLOT
- * \retval #PSA_ERROR_NOT_PERMITTED
- * \retval #PSA_ERROR_INVALID_ARGUMENT
- *         \c key is not compatible with \c alg,
- *         or \p capacity is too large for the specified algorithm and key.
- * \retval #PSA_ERROR_NOT_SUPPORTED
- *         \c alg is not supported or is not a key derivation algorithm.
- * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
- * \retval #PSA_ERROR_COMMUNICATION_FAILURE
- * \retval #PSA_ERROR_HARDWARE_FAILURE
- * \retval #PSA_ERROR_TAMPERING_DETECTED
- * \retval #PSA_ERROR_BAD_STATE
- *         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 psa_key_derivation(psa_crypto_generator_t *generator,
-                                psa_key_handle_t handle,
-                                psa_algorithm_t alg,
-                                const uint8_t *salt,
-                                size_t salt_length,
-                                const uint8_t *label,
-                                size_t label_length,
-                                size_t capacity);
-
 /** Perform a key agreement and use the shared secret as input to a key
  * derivation.
  *
diff --git a/include/psa/crypto_extra.h b/include/psa/crypto_extra.h
index 7f0885794..7d89fbfe0 100644
--- a/include/psa/crypto_extra.h
+++ b/include/psa/crypto_extra.h
@@ -120,6 +120,62 @@ void mbedtls_psa_crypto_free( void );
 psa_status_t mbedtls_psa_inject_entropy(const unsigned char *seed,
                                         size_t seed_size);
 
+/** Set up a key derivation operation.
+ *
+ * FIMXE This function is no longer part of the official API. Its prototype
+ * is only kept around for the sake of tests that haven't been updated yet.
+ *
+ * A key derivation algorithm takes three inputs: a secret input \p key and
+ * two non-secret inputs \p label and p salt.
+ * The result of this function is a byte generator which can
+ * be used to produce keys and other cryptographic material.
+ *
+ * The role of \p label and \p salt is as follows:
+ * - For HKDF (#PSA_ALG_HKDF), \p salt is the salt used in the "extract" step
+ *   and \p label is the info string used in the "expand" step.
+ *
+ * \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.
+ * \param handle                  Handle to the secret key.
+ * \param alg                     The key derivation algorithm to compute
+ *                                (\c PSA_ALG_XXX value such that
+ *                                #PSA_ALG_IS_KEY_DERIVATION(\p alg) is true).
+ * \param[in] salt                Salt to use.
+ * \param salt_length             Size of the \p salt buffer in bytes.
+ * \param[in] label               Label to use.
+ * \param label_length            Size of the \p label buffer in bytes.
+ * \param capacity                The maximum number of bytes that the
+ *                                generator will be able to provide.
+ *
+ * \retval #PSA_SUCCESS
+ *         Success.
+ * \retval #PSA_ERROR_INVALID_HANDLE
+ * \retval #PSA_ERROR_EMPTY_SLOT
+ * \retval #PSA_ERROR_NOT_PERMITTED
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         \c key is not compatible with \c alg,
+ *         or \p capacity is too large for the specified algorithm and key.
+ * \retval #PSA_ERROR_NOT_SUPPORTED
+ *         \c alg is not supported or is not a key derivation algorithm.
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ * \retval #PSA_ERROR_COMMUNICATION_FAILURE
+ * \retval #PSA_ERROR_HARDWARE_FAILURE
+ * \retval #PSA_ERROR_TAMPERING_DETECTED
+ * \retval #PSA_ERROR_BAD_STATE
+ *         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 psa_key_derivation(psa_crypto_generator_t *generator,
+                                psa_key_handle_t handle,
+                                psa_algorithm_t alg,
+                                const uint8_t *salt,
+                                size_t salt_length,
+                                const uint8_t *label,
+                                size_t label_length,
+                                size_t capacity);
+
 
 #ifdef __cplusplus
 }