diff --git a/include/psa/crypto_extra.h b/include/psa/crypto_extra.h
index 6de464d..eaadf81 100644
--- a/include/psa/crypto_extra.h
+++ b/include/psa/crypto_extra.h
@@ -1760,7 +1760,13 @@
  *                      recognized, or the parameters are incompatible,
  *                      return 0.
  */
-#define PSA_PAKE_OUTPUT_SIZE(alg, primitive, output_step) 0
+#define PSA_PAKE_OUTPUT_SIZE(alg, primitive, output_step)                    \
+            ( alg == PSA_ALG_JPAKE &&                                        \
+              primitive == PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC,   \
+                                              PSA_ECC_FAMILY_SECP_R1, 256) ? \
+              ( output_step == PSA_PAKE_STEP_KEY_SHARE ? 69 :                \
+                ( output_step == PSA_PAKE_STEP_ZK_PUBLIC ? 66 : 33 ) ) : 0 )
+
 
 /** A sufficient input buffer size for psa_pake_input().
  *
@@ -1781,7 +1787,12 @@
  *                      the input type or PAKE algorithm is not recognized, or
  *                      the parameters are incompatible, return 0.
  */
-#define PSA_PAKE_INPUT_SIZE(alg, primitive, input_step) 0
+#define PSA_PAKE_INPUT_SIZE(alg, primitive, input_step)                      \
+            ( alg == PSA_ALG_JPAKE &&                                        \
+              primitive == PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC,   \
+                                              PSA_ECC_FAMILY_SECP_R1, 256) ? \
+              ( input_step == PSA_PAKE_STEP_KEY_SHARE ? 69 :                 \
+                ( input_step == PSA_PAKE_STEP_ZK_PUBLIC ? 66 : 33 ) ) : 0 )
 
 /** Output buffer size for psa_pake_output() for any of the supported PAKE
  * algorithm and primitive suites and output step.
@@ -1790,7 +1801,7 @@
  *
  * See also #PSA_PAKE_OUTPUT_SIZE(\p alg, \p primitive, \p step).
  */
-#define PSA_PAKE_OUTPUT_MAX_SIZE 0
+#define PSA_PAKE_OUTPUT_MAX_SIZE 69
 
 /** Input buffer size for psa_pake_input() for any of the supported PAKE
  * algorithm and primitive suites and input step.
@@ -1799,7 +1810,7 @@
  *
  * See also #PSA_PAKE_INPUT_SIZE(\p alg, \p primitive, \p step).
  */
-#define PSA_PAKE_INPUT_MAX_SIZE 0
+#define PSA_PAKE_INPUT_MAX_SIZE 69
 
 /** Returns a suitable initializer for a PAKE cipher suite object of type
  * psa_pake_cipher_suite_t.
@@ -1809,7 +1820,11 @@
 /** Returns a suitable initializer for a PAKE operation object of type
  * psa_pake_operation_t.
  */
-#define PSA_PAKE_OPERATION_INIT {PSA_ALG_NONE, {0}}
+#if defined(MBEDTLS_PSA_BUILTIN_PAKE)
+#define PSA_PAKE_OPERATION_INIT {PSA_ALG_NONE, 0, 0, 0, 0, MBEDTLS_SVC_KEY_ID_INIT, 0, NULL, 0, 0, { .dummy = 0 } }
+#else
+#define PSA_PAKE_OPERATION_INIT {PSA_ALG_NONE, 0, 0, {0}}
+#endif
 
 struct psa_pake_cipher_suite_s
 {
@@ -1879,14 +1894,32 @@
         cipher_suite->hash = hash;
 }
 
+#if defined(MBEDTLS_PSA_BUILTIN_PAKE)
+#include <mbedtls/ecjpake.h>
+#endif
+
 struct psa_pake_operation_s
 {
-    psa_algorithm_t alg;
+    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);
+    mbedtls_svc_key_id_t MBEDTLS_PRIVATE(password);
+    psa_pake_role_t MBEDTLS_PRIVATE(role);
+    uint8_t *MBEDTLS_PRIVATE(buffer);
+    size_t MBEDTLS_PRIVATE(buffer_length);
+    size_t MBEDTLS_PRIVATE(buffer_offset);
+#endif
     union
     {
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_ECJPAKE)
+        mbedtls_ecjpake_context ecjpake;
+#endif
         /* Make the union non-empty even with no supported algorithms. */
         uint8_t dummy;
-    } ctx;
+    } MBEDTLS_PRIVATE(ctx);
 };
 
 static inline struct psa_pake_cipher_suite_s psa_pake_cipher_suite_init( void )