diff --git a/library/pk_wrap.h b/library/pk_wrap.h
index eead322..ae970f9 100644
--- a/library/pk_wrap.h
+++ b/library/pk_wrap.h
@@ -27,6 +27,10 @@
 
 #include "mbedtls/pk.h"
 
+#if defined(MBEDTLS_PSA_CRYPTO_C)
+#include "psa/crypto.h"
+#endif /* MBEDTLS_PSA_CRYPTO_C */
+
 struct mbedtls_pk_info_t
 {
     /** Public key type */
@@ -135,7 +139,7 @@
 extern const mbedtls_pk_info_t mbedtls_pk_opaque_info;
 #endif
 
-#if defined(MBEDTLS_USE_PSA_CRYPTO)
+#if defined(MBEDTLS_PSA_CRYPTO_C)
 int mbedtls_pk_error_from_psa( psa_status_t status );
 
 #if defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY)
@@ -146,9 +150,6 @@
 int mbedtls_pk_error_from_psa_rsa( psa_status_t status );
 #endif
 
-#endif /* MBEDTLS_USE_PSA_CRYPTO */
-
-#if defined(MBEDTLS_USE_PSA_CRYPTO)
 int mbedtls_pk_psa_sign_ext( psa_algorithm_t psa_alg_md, void *ctx,
                              const unsigned char *hash, size_t hash_len,
                              unsigned char *sig, size_t sig_size,