diff --git a/include/psa/crypto_extra.h b/include/psa/crypto_extra.h
index 4920508d7..0e2d57ce2 100644
--- a/include/psa/crypto_extra.h
+++ b/include/psa/crypto_extra.h
@@ -1960,7 +1960,7 @@ psa_status_t psa_pake_abort(psa_pake_operation_t *operation);
 /** Returns a suitable initializer for a PAKE operation object of type
  * psa_pake_operation_t.
  */
-#define PSA_PAKE_OPERATION_INIT { 0, PSA_ALG_NONE, PSA_PAKE_OPERATION_STAGE_SETUP, \
+#define PSA_PAKE_OPERATION_INIT { 0, PSA_ALG_NONE, 0, PSA_PAKE_OPERATION_STAGE_SETUP, \
                                   { 0 }, { { 0 } } }
 
 struct psa_pake_cipher_suite_s {
@@ -2106,6 +2106,8 @@ struct psa_pake_operation_s {
     unsigned int MBEDTLS_PRIVATE(id);
     /* Algorithm of the PAKE operation */
     psa_algorithm_t MBEDTLS_PRIVATE(alg);
+    /* A primitive of type compatible with algorithm */
+    psa_pake_primitive_t MBEDTLS_PRIVATE(primitive);
     /* Stage of the PAKE operation: waiting for the setup, collecting inputs
      * or computing. */
     uint8_t MBEDTLS_PRIVATE(stage);
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index bc19ed07c..049edfcac 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -7323,6 +7323,8 @@ psa_status_t psa_pake_setup(
     memset(&operation->data.inputs, 0, sizeof(operation->data.inputs));
 
     operation->alg = cipher_suite->algorithm;
+    operation->primitive = PSA_PAKE_PRIMITIVE(cipher_suite->type,
+                                              cipher_suite->family, cipher_suite->bits);
     operation->data.inputs.cipher_suite = *cipher_suite;
 
 #if defined(PSA_WANT_ALG_JPAKE)
@@ -7920,7 +7922,9 @@ psa_status_t psa_pake_input(
         goto exit;
     }
 
-    if (input_length == 0 || input_length > PSA_PAKE_INPUT_MAX_SIZE) {
+    if (input_length == 0 || input_length > PSA_PAKE_INPUT_SIZE(operation->alg,
+                                                                operation->primitive,
+                                                                step)) {
         status = PSA_ERROR_INVALID_ARGUMENT;
         goto exit;
     }