diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 67fdd5c..0f7d4bd 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -6013,7 +6013,7 @@
         *key_buffer_size = PSA_BITS_TO_BYTES( bits );
     }
     else
-#if defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR)
+#if defined(PSA_WANT_KEY_TYPE_RSA_KEY_PAIR)
     if( PSA_KEY_TYPE_IS_RSA( type ) && PSA_KEY_TYPE_IS_KEY_PAIR( type ) )
     {
         if( bits > PSA_VENDOR_RSA_MAX_KEY_BITS )
@@ -6027,15 +6027,15 @@
         *key_buffer_size = PSA_KEY_EXPORT_RSA_KEY_PAIR_MAX_SIZE( bits );
     }
     else
-#endif /* defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR) */
+#endif /* defined(PSA_WANT_KEY_TYPE_RSA_KEY_PAIR) */
 
-#if defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_KEY_PAIR)
+#if defined(PSA_WANT_KEY_TYPE_ECC_KEY_PAIR)
     if( PSA_KEY_TYPE_IS_ECC( type ) && PSA_KEY_TYPE_IS_KEY_PAIR( type ) )
     {
         *key_buffer_size = PSA_BITS_TO_BYTES( bits );
     }
     else
-#endif /* defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_KEY_PAIR) */
+#endif /* defined(PSA_WANT_KEY_TYPE_ECC_KEY_PAIR) */
     {
         return( PSA_ERROR_NOT_SUPPORTED );
     }