diff --git a/include/psa/crypto_builtin_composites.h b/include/psa/crypto_builtin_composites.h
index 9f23551eb1..295452c8cc 100644
--- a/include/psa/crypto_builtin_composites.h
+++ b/include/psa/crypto_builtin_composites.h
@@ -180,5 +180,41 @@ typedef struct {
 #endif
 
 
+/* EC-JPAKE operation definitions */
+
+#include "mbedtls/ecjpake.h"
+
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
+#define MBEDTLS_PSA_BUILTIN_PAKE  1
+#endif
+
+/* 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)
+
+typedef struct {
+    psa_algorithm_t MBEDTLS_PRIVATE(alg);
+    unsigned int MBEDTLS_PRIVATE(state);
+    unsigned int MBEDTLS_PRIVATE(sequence);
+#if defined(MBEDTLS_PSA_BUILTIN_PAKE)
+    unsigned int MBEDTLS_PRIVATE(input_step);
+    unsigned int MBEDTLS_PRIVATE(output_step);
+    uint8_t *MBEDTLS_PRIVATE(password);
+    size_t MBEDTLS_PRIVATE(password_len);
+    uint8_t MBEDTLS_PRIVATE(role);
+    uint8_t MBEDTLS_PRIVATE(buffer[MBEDTLS_PSA_PAKE_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);
+        mbedtls_ecjpake_context MBEDTLS_PRIVATE(pake);
+    } MBEDTLS_PRIVATE(ctx);
+
+} mbedtls_psa_pake_operation_t;
+
+#define MBEDTLS_PSA_PAKE_OPERATION_INIT { { 0 } }
 
 #endif /* PSA_CRYPTO_BUILTIN_COMPOSITES_H */
diff --git a/include/psa/crypto_builtin_primitives.h b/include/psa/crypto_builtin_primitives.h
index 2830b61e6c..c76bc78143 100644
--- a/include/psa/crypto_builtin_primitives.h
+++ b/include/psa/crypto_builtin_primitives.h
@@ -111,42 +111,4 @@ typedef struct {
 
 #define MBEDTLS_PSA_CIPHER_OPERATION_INIT { 0, 0, 0, { 0 } }
 
-
-/* EC-JPAKE operation definitions */
-
-#include "mbedtls/ecjpake.h"
-
-#if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
-#define MBEDTLS_PSA_BUILTIN_PAKE  1
-#endif
-
-/* 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)
-
-typedef struct {
-    psa_algorithm_t MBEDTLS_PRIVATE(alg);
-    unsigned int MBEDTLS_PRIVATE(state);
-    unsigned int MBEDTLS_PRIVATE(sequence);
-#if defined(MBEDTLS_PSA_BUILTIN_PAKE)
-    unsigned int MBEDTLS_PRIVATE(input_step);
-    unsigned int MBEDTLS_PRIVATE(output_step);
-    uint8_t *MBEDTLS_PRIVATE(password);
-    size_t MBEDTLS_PRIVATE(password_len);
-    uint8_t MBEDTLS_PRIVATE(role);
-    uint8_t MBEDTLS_PRIVATE(buffer[MBEDTLS_PSA_PAKE_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);
-        mbedtls_ecjpake_context MBEDTLS_PRIVATE(pake);
-    } MBEDTLS_PRIVATE(ctx);
-
-} mbedtls_psa_pake_operation_t;
-
-#define MBEDTLS_PSA_PAKE_OPERATION_INIT { { 0 } }
-
 #endif /* PSA_CRYPTO_BUILTIN_PRIMITIVES_H */
diff --git a/include/psa/crypto_driver_contexts_composites.h b/include/psa/crypto_driver_contexts_composites.h
index 1b95814f9e..4d0e9848d3 100644
--- a/include/psa/crypto_driver_contexts_composites.h
+++ b/include/psa/crypto_driver_contexts_composites.h
@@ -88,6 +88,32 @@ typedef mbedtls_psa_aead_operation_t
 
 #endif /* MBEDTLS_TEST_LIBTESTDRIVER1 && LIBTESTDRIVER1_MBEDTLS_PSA_BUILTIN_AEAD */
 
+#if defined(MBEDTLS_TEST_LIBTESTDRIVER1) && \
+    defined(LIBTESTDRIVER1_MBEDTLS_PSA_BUILTIN_PAKE)
+
+typedef libtestdriver1_mbedtls_psa_pake_operation_t
+    mbedtls_transparent_test_driver_pake_operation_t;
+typedef libtestdriver1_psa_pake_operation_t
+    mbedtls_opaque_test_driver_pake_operation_t;
+
+#define MBEDTLS_TRANSPARENT_TEST_DRIVER_PAKE_OPERATION_INIT \
+    LIBTESTDRIVER1_MBEDTLS_PSA_PAKE_OPERATION_INIT
+#define MBEDTLS_OPAQUE_TEST_DRIVER_PAKE_OPERATION_INIT \
+    LIBTESTDRIVER1_MBEDTLS_PSA_PAKE_OPERATION_INIT
+
+#else
+typedef mbedtls_psa_pake_operation_t
+    mbedtls_transparent_test_driver_pake_operation_t;
+typedef mbedtls_psa_pake_operation_t
+    mbedtls_opaque_test_driver_pake_operation_t;
+
+#define MBEDTLS_TRANSPARENT_TEST_DRIVER_PAKE_OPERATION_INIT \
+    MBEDTLS_PSA_PAKE_OPERATION_INIT
+#define MBEDTLS_OPAQUE_TEST_DRIVER_PAKE_OPERATION_INIT \
+    MBEDTLS_PSA_PAKE_OPERATION_INIT
+
+#endif /* MBEDTLS_TEST_LIBTESTDRIVER1 && LIBTESTDRIVER1_MBEDTLS_PSA_BUILTIN_PAKE */
+
 #endif /* PSA_CRYPTO_DRIVER_TEST */
 
 /* Define the context to be used for an operation that is executed through the
@@ -124,5 +150,14 @@ typedef union {
     mbedtls_psa_verify_hash_interruptible_operation_t mbedtls_ctx;
 } psa_driver_verify_hash_interruptible_context_t;
 
+typedef union {
+    unsigned dummy; /* Make sure this union is always non-empty */
+    mbedtls_psa_pake_operation_t mbedtls_ctx;
+#if defined(PSA_CRYPTO_DRIVER_TEST)
+    mbedtls_transparent_test_driver_pake_operation_t transparent_test_driver_ctx;
+    mbedtls_opaque_test_driver_pake_operation_t opaque_test_driver_ctx;
+#endif
+} psa_driver_pake_context_t;
+
 #endif /* PSA_CRYPTO_DRIVER_CONTEXTS_COMPOSITES_H */
 /* End of automatically generated file. */
diff --git a/include/psa/crypto_driver_contexts_primitives.h b/include/psa/crypto_driver_contexts_primitives.h
index 7d096208b4..f1463f34d0 100644
--- a/include/psa/crypto_driver_contexts_primitives.h
+++ b/include/psa/crypto_driver_contexts_primitives.h
@@ -89,32 +89,6 @@ typedef struct {
 #define MBEDTLS_OPAQUE_TEST_DRIVER_CIPHER_OPERATION_INIT \
     { 0, MBEDTLS_TRANSPARENT_TEST_DRIVER_CIPHER_OPERATION_INIT }
 
-#if defined(MBEDTLS_TEST_LIBTESTDRIVER1) && \
-    defined(LIBTESTDRIVER1_MBEDTLS_PSA_BUILTIN_PAKE)
-
-typedef libtestdriver1_mbedtls_psa_pake_operation_t
-    mbedtls_transparent_test_driver_pake_operation_t;
-typedef libtestdriver1_psa_pake_operation_t
-    mbedtls_opaque_test_driver_pake_operation_t;
-
-#define MBEDTLS_TRANSPARENT_TEST_DRIVER_PAKE_OPERATION_INIT \
-    LIBTESTDRIVER1_MBEDTLS_PSA_PAKE_OPERATION_INIT
-#define MBEDTLS_OPAQUE_TEST_DRIVER_PAKE_OPERATION_INIT \
-    LIBTESTDRIVER1_MBEDTLS_PSA_PAKE_OPERATION_INIT
-
-#else
-typedef mbedtls_psa_pake_operation_t
-    mbedtls_transparent_test_driver_pake_operation_t;
-typedef mbedtls_psa_pake_operation_t
-    mbedtls_opaque_test_driver_pake_operation_t;
-
-#define MBEDTLS_TRANSPARENT_TEST_DRIVER_PAKE_OPERATION_INIT \
-    MBEDTLS_PSA_PAKE_OPERATION_INIT
-#define MBEDTLS_OPAQUE_TEST_DRIVER_PAKE_OPERATION_INIT \
-    MBEDTLS_PSA_PAKE_OPERATION_INIT
-
-#endif /* MBEDTLS_TEST_LIBTESTDRIVER1 && LIBTESTDRIVER1_MBEDTLS_PSA_BUILTIN_PAKE */
-
 #endif /* PSA_CRYPTO_DRIVER_TEST */
 
 /* Define the context to be used for an operation that is executed through the
@@ -141,14 +115,5 @@ typedef union {
 #endif
 } psa_driver_cipher_context_t;
 
-typedef union {
-    unsigned dummy; /* Make sure this union is always non-empty */
-    mbedtls_psa_pake_operation_t mbedtls_ctx;
-#if defined(PSA_CRYPTO_DRIVER_TEST)
-    mbedtls_transparent_test_driver_pake_operation_t transparent_test_driver_ctx;
-    mbedtls_opaque_test_driver_pake_operation_t opaque_test_driver_ctx;
-#endif
-} psa_driver_pake_context_t;
-
 #endif /* PSA_CRYPTO_DRIVER_CONTEXTS_PRIMITIVES_H */
 /* End of automatically generated file. */