diff --git a/programs/psa/crypto_examples.c b/programs/psa/crypto_examples.c
index 935d657..e9323a5 100644
--- a/programs/psa/crypto_examples.c
+++ b/programs/psa/crypto_examples.c
@@ -43,13 +43,13 @@
         }                                                                     \
     } while ( 0 )
 
-#if !defined(MBEDTLS_PSA_CRYPTO_C) || !defined(MBEDTLS_AES_C) || \
+#if !defined(MBEDTLS_PSA_CRYPTO_C) || !defined(PSA_WANT_KEY_TYPE_AES) || \
     !defined(MBEDTLS_CIPHER_MODE_CBC) || !defined(MBEDTLS_CIPHER_MODE_CTR) || \
     !defined(MBEDTLS_CIPHER_MODE_WITH_PADDING) || \
     defined(MBEDTLS_PSA_CRYPTO_KEY_ID_ENCODES_OWNER)
 int main( void )
 {
-    printf( "MBEDTLS_PSA_CRYPTO_C and/or MBEDTLS_AES_C and/or "
+    printf( "MBEDTLS_PSA_CRYPTO_C and/or PSA_WANT_KEY_TYPE_AES and/or "
             "MBEDTLS_CIPHER_MODE_CBC and/or MBEDTLS_CIPHER_MODE_CTR "
             "and/or MBEDTLS_CIPHER_MODE_WITH_PADDING "
             "not defined and/or MBEDTLS_PSA_CRYPTO_KEY_ID_ENCODES_OWNER"
@@ -327,5 +327,5 @@
     mbedtls_psa_crypto_free( );
     return( 0 );
 }
-#endif /* MBEDTLS_PSA_CRYPTO_C && MBEDTLS_AES_C && MBEDTLS_CIPHER_MODE_CBC &&
+#endif /* MBEDTLS_PSA_CRYPTO_C && PSA_WANT_KEY_TYPE_AES && MBEDTLS_CIPHER_MODE_CBC &&
           MBEDTLS_CIPHER_MODE_CTR && MBEDTLS_CIPHER_MODE_WITH_PADDING */
diff --git a/programs/psa/key_ladder_demo.c b/programs/psa/key_ladder_demo.c
index 47d5de6..2abb1b8 100644
--- a/programs/psa/key_ladder_demo.c
+++ b/programs/psa/key_ladder_demo.c
@@ -66,13 +66,13 @@
 
 /* If the build options we need are not enabled, compile a placeholder. */
 #if !defined(MBEDTLS_SHA256_C) || !defined(MBEDTLS_MD_C) ||      \
-    !defined(MBEDTLS_AES_C) || !defined(MBEDTLS_CCM_C) ||        \
+    !defined(PSA_WANT_KEY_TYPE_AES) || !defined(MBEDTLS_CCM_C) ||        \
     !defined(MBEDTLS_PSA_CRYPTO_C) || !defined(MBEDTLS_FS_IO) || \
     defined(MBEDTLS_PSA_CRYPTO_KEY_ID_ENCODES_OWNER)
 int main( void )
 {
     printf( "MBEDTLS_SHA256_C and/or MBEDTLS_MD_C and/or "
-            "MBEDTLS_AES_C and/or MBEDTLS_CCM_C and/or "
+            "PSA_WANT_KEY_TYPE_AES and/or MBEDTLS_CCM_C and/or "
             "MBEDTLS_PSA_CRYPTO_C and/or MBEDTLS_FS_IO "
             "not defined and/or MBEDTLS_PSA_CRYPTO_KEY_ID_ENCODES_OWNER "
             "defined.\n" );
@@ -695,4 +695,4 @@
     usage( );
     return( EXIT_FAILURE );
 }
-#endif /* MBEDTLS_SHA256_C && MBEDTLS_MD_C && MBEDTLS_AES_C && MBEDTLS_CCM_C && MBEDTLS_PSA_CRYPTO_C && MBEDTLS_FS_IO */
+#endif /* MBEDTLS_SHA256_C && MBEDTLS_MD_C && PSA_WANT_KEY_TYPE_AES && MBEDTLS_CCM_C && MBEDTLS_PSA_CRYPTO_C && MBEDTLS_FS_IO */