diff --git a/include/mbedtls/pk.h b/include/mbedtls/pk.h
index 64cd7b5..09a8513 100644
--- a/include/mbedtls/pk.h
+++ b/include/mbedtls/pk.h
@@ -142,7 +142,9 @@
  */
 typedef struct mbedtls_pk_context
 {
+#if !defined(MBEDTLS_PK_SINGLE_TYPE)
     mbedtls_pk_handle_t         pk_info; /**< Public key information         */
+#endif
     void *                      pk_ctx;  /**< Underlying public key context  */
 } mbedtls_pk_context;
 
diff --git a/include/mbedtls/pk_internal.h b/include/mbedtls/pk_internal.h
index 94dd9aa..5258cb1 100644
--- a/include/mbedtls/pk_internal.h
+++ b/include/mbedtls/pk_internal.h
@@ -234,7 +234,11 @@
 /*
  * Macros to access pk_info
  */
+#if defined(MBEDTLS_PK_SINGLE_TYPE)
+#define MBEDTLS_PK_CTX_INFO( ctx )      MBEDTLS_PK_UNIQUE_VALID_HANDLE
+#else
 #define MBEDTLS_PK_CTX_INFO( ctx )      ( (ctx)->pk_info )
+#endif
 #define MBEDTLS_PK_CTX_IS_VALID( ctx )  \
     ( MBEDTLS_PK_CTX_INFO( (ctx) ) != MBEDTLS_PK_INVALID_HANDLE )
 
diff --git a/library/pk.c b/library/pk.c
index d6bf976..dfa9a68 100644
--- a/library/pk.c
+++ b/library/pk.c
@@ -1299,7 +1299,9 @@
 {
     PK_VALIDATE( ctx != NULL );
 
+#if !defined(MBEDTLS_PK_SINGLE_TYPE)
     ctx->pk_info = MBEDTLS_PK_INVALID_HANDLE;
+#endif
     ctx->pk_ctx = NULL;
 }
 
@@ -1394,17 +1396,19 @@
 int mbedtls_pk_setup( mbedtls_pk_context *ctx, mbedtls_pk_handle_t info )
 {
     PK_VALIDATE_RET( ctx != NULL );
-    if( info == MBEDTLS_PK_INVALID_HANDLE ||
-        MBEDTLS_PK_CTX_IS_VALID( ctx ) )
-    {
+    if( info == MBEDTLS_PK_INVALID_HANDLE )
         return( MBEDTLS_ERR_PK_BAD_INPUT_DATA );
-    }
+
+#if !defined(MBEDTLS_PK_SINGLE_TYPE)
+    if( ctx->pk_info != NULL )
+        return( MBEDTLS_ERR_PK_BAD_INPUT_DATA );
+
+    ctx->pk_info = info;
+#endif
 
     if( ( ctx->pk_ctx = pk_info_ctx_alloc_func( info ) ) == NULL )
         return( MBEDTLS_ERR_PK_ALLOC_FAILED );
 
-    ctx->pk_info = info;
-
     return( 0 );
 }
 
diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index 1635855..b446f49 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -95,9 +95,11 @@
 void valid_parameters( )
 {
     mbedtls_pk_context pk;
+#if !defined(MBEDTLS_PK_SINGLE_TYPE)
     unsigned char buf[1];
     size_t len;
     void *options = NULL;
+#endif
 
     mbedtls_pk_init( &pk );
 
@@ -118,6 +120,7 @@
     TEST_ASSERT( mbedtls_pk_get_len( NULL ) == 0 );
     TEST_ASSERT( mbedtls_pk_can_do( NULL, MBEDTLS_PK_NONE ) == 0 );
 
+#if !defined(MBEDTLS_PK_SINGLE_TYPE)
     TEST_ASSERT( mbedtls_pk_sign_restartable( &pk,
                                               MBEDTLS_MD_NONE,
                                               NULL, 0,
@@ -172,6 +175,7 @@
                                      NULL, &len, 0,
                                      rnd_std_rand, NULL ) ==
                  MBEDTLS_ERR_PK_BAD_INPUT_DATA );
+#endif /* MBEDTLS_PK_SINGLE_TYPE */
 
 #if defined(MBEDTLS_PK_PARSE_C)
     TEST_ASSERT( mbedtls_pk_parse_key( &pk, NULL, 0, NULL, 1 ) ==