diff --git a/include/psa/crypto_driver_contexts_key_derivation.h b/include/psa/crypto_driver_contexts_key_derivation.h
index 32de4f765..3fb29ff7f 100644
--- a/include/psa/crypto_driver_contexts_key_derivation.h
+++ b/include/psa/crypto_driver_contexts_key_derivation.h
@@ -55,8 +55,7 @@ typedef union {
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_TLS12_ECJPAKE_TO_PMS)
     psa_tls12_ecjpake_to_pms_t MBEDTLS_PRIVATE(tls12_ecjpake_to_pms);
 #endif
-#if defined(MBEDTLS_PSA_BUILTIN_ALG_PBKDF2_HMAC) || \
-    defined(MBEDTLS_PSA_BUILTIN_ALG_PBKDF2_AES_CMAC_PRF_128)
+#if defined(PSA_HAVE_SOFT_PBKDF2)
     psa_pbkdf2_key_derivation_t MBEDTLS_PRIVATE(pbkdf2);
 #endif
 } psa_driver_key_derivation_context_t;