diff --git a/include/psa/crypto_builtin_composites.h b/include/psa/crypto_builtin_composites.h
index 3221a64..f331ec5 100644
--- a/include/psa/crypto_builtin_composites.h
+++ b/include/psa/crypto_builtin_composites.h
@@ -191,23 +191,25 @@
 /* Note: the format for mbedtls_ecjpake_read/write function has an extra
  * length byte for each step, plus an extra 3 bytes for ECParameters in the
  * server's 2nd round. */
-#define MBEDTLS_PSA_PAKE_BUFFER_SIZE ((3 + 1 + 65 + 1 + 65 + 1 + 32) * 2)
+#define MBEDTLS_PSA_JPAKE_BUFFER_SIZE ((3 + 1 + 65 + 1 + 65 + 1 + 32) * 2)
 
 typedef struct {
     psa_algorithm_t MBEDTLS_PRIVATE(alg);
 
-#if defined(MBEDTLS_PSA_BUILTIN_PAKE)
     uint8_t *MBEDTLS_PRIVATE(password);
     size_t MBEDTLS_PRIVATE(password_len);
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
     uint8_t MBEDTLS_PRIVATE(role);
-    uint8_t MBEDTLS_PRIVATE(buffer[MBEDTLS_PSA_PAKE_BUFFER_SIZE]);
+    uint8_t MBEDTLS_PRIVATE(buffer[MBEDTLS_PSA_JPAKE_BUFFER_SIZE]);
     size_t MBEDTLS_PRIVATE(buffer_length);
     size_t MBEDTLS_PRIVATE(buffer_offset);
 #endif
     /* Context structure for the Mbed TLS EC-JPAKE implementation. */
     union {
         unsigned int MBEDTLS_PRIVATE(dummy);
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
         mbedtls_ecjpake_context MBEDTLS_PRIVATE(pake);
+#endif
     } MBEDTLS_PRIVATE(ctx);
 
 } mbedtls_psa_pake_operation_t;
diff --git a/include/psa/crypto_driver_contexts_composites.h b/include/psa/crypto_driver_contexts_composites.h
index 4d0e984..6c56a51 100644
--- a/include/psa/crypto_driver_contexts_composites.h
+++ b/include/psa/crypto_driver_contexts_composites.h
@@ -93,7 +93,7 @@
 
 typedef libtestdriver1_mbedtls_psa_pake_operation_t
     mbedtls_transparent_test_driver_pake_operation_t;
-typedef libtestdriver1_psa_pake_operation_t
+typedef libtestdriver1_mbedtls_psa_pake_operation_t
     mbedtls_opaque_test_driver_pake_operation_t;
 
 #define MBEDTLS_TRANSPARENT_TEST_DRIVER_PAKE_OPERATION_INIT \
diff --git a/include/psa/crypto_driver_contexts_primitives.h b/include/psa/crypto_driver_contexts_primitives.h
index f1463f3..620a4b3 100644
--- a/include/psa/crypto_driver_contexts_primitives.h
+++ b/include/psa/crypto_driver_contexts_primitives.h
@@ -45,8 +45,6 @@
 #include <libtestdriver1/include/psa/crypto.h>
 #endif
 
-#include "mbedtls/ecjpake.h"
-
 #if defined(PSA_CRYPTO_DRIVER_TEST)
 
 #if defined(MBEDTLS_TEST_LIBTESTDRIVER1) && \
diff --git a/include/psa/crypto_extra.h b/include/psa/crypto_extra.h
index 8b8cb04..39ef52c 100644
--- a/include/psa/crypto_extra.h
+++ b/include/psa/crypto_extra.h
@@ -429,7 +429,7 @@
  */
 #define PSA_DH_FAMILY_CUSTOM             ((psa_dh_family_t) 0x7e)
 
-/** EC-JPAKE operation stages. */
+/** PAKE operation stages. */
 #define PSA_PAKE_OPERATION_STAGE_SETUP 0
 #define PSA_PAKE_OPERATION_STAGE_COLLECT_INPUTS 1
 #define PSA_PAKE_OPERATION_STAGE_COMPUTATION 2
@@ -1895,7 +1895,7 @@
  * psa_pake_operation_t.
  */
 #define PSA_PAKE_OPERATION_INIT { 0, PSA_ALG_NONE, PSA_PAKE_OPERATION_STAGE_SETUP, \
-                                  { { 0, 0, 0, 0 } }, { { 0 } } }
+                                  { 0 }, { { 0 } } }
 
 struct psa_pake_cipher_suite_s {
     psa_algorithm_t algorithm;
@@ -2002,7 +2002,7 @@
     PSA_PAKE_SEQ_END            = 7,
 };
 
-typedef enum psa_pake_driver_step {
+typedef enum psa_crypto_driver_pake_step {
     PSA_JPAKE_STEP_INVALID        = 0,  /* Invalid step */
     PSA_JPAKE_X1_STEP_KEY_SHARE   = 1,  /* Round 1: input/output key share (for ephemeral private key X1).*/
     PSA_JPAKE_X1_STEP_ZK_PUBLIC   = 2,  /* Round 1: input/output Schnorr NIZKP public key for the X1 key */
@@ -2016,7 +2016,7 @@
     PSA_JPAKE_X4S_STEP_KEY_SHARE  = 10, /* Round 2: input X4S key (from peer) */
     PSA_JPAKE_X4S_STEP_ZK_PUBLIC  = 11, /* Round 2: input Schnorr NIZKP public key for the X4S key (from peer) */
     PSA_JPAKE_X4S_STEP_ZK_PROOF   = 12  /* Round 2: input Schnorr NIZKP proof for the X4S key (from peer) */
-} psa_pake_driver_step_t;
+} psa_crypto_driver_pake_step_t;
 
 
 struct psa_jpake_computation_stage_s {
@@ -2042,7 +2042,10 @@
     uint8_t MBEDTLS_PRIVATE(stage);
     /* Holds computation stage of the PAKE algorithms. */
     union {
+        uint8_t MBEDTLS_PRIVATE(dummy);
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
         psa_jpake_computation_stage_t MBEDTLS_PRIVATE(jpake);
+#endif
     } MBEDTLS_PRIVATE(computation_stage);
     union {
         psa_driver_pake_context_t MBEDTLS_PRIVATE(ctx);