diff --git a/include/psa/crypto_extra.h b/include/psa/crypto_extra.h
index 07d7bae31..32e956925 100644
--- a/include/psa/crypto_extra.h
+++ b/include/psa/crypto_extra.h
@@ -430,8 +430,9 @@ psa_status_t mbedtls_psa_inject_entropy(const uint8_t *seed,
 #define PSA_DH_FAMILY_CUSTOM             ((psa_dh_family_t) 0x7e)
 
 /** EC-JPAKE operation stages. */
-#define PSA_PAKE_OPERATION_STAGE_COLLECT_INPUTS 0
-#define PSA_PAKE_OPERATION_STAGE_COMPUTATION 1
+#define PSA_PAKE_OPERATION_STAGE_SETUP 0
+#define PSA_PAKE_OPERATION_STAGE_COLLECT_INPUTS 1
+#define PSA_PAKE_OPERATION_STAGE_COMPUTATION 2
 
 /**
  * \brief Set domain parameters for a key.
@@ -1893,7 +1894,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_COLLECT_INPUTS, \
+#define PSA_PAKE_OPERATION_INIT { 0, PSA_ALG_NONE, PSA_PAKE_OPERATION_STAGE_SETUP, \
                                   { { 0 } }, { { 0 } } }
 
 struct psa_pake_cipher_suite_s {
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index bfbd49735..2d1c06500 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -7237,11 +7237,7 @@ psa_status_t psa_pake_setup(
     psa_pake_operation_t *operation,
     const psa_pake_cipher_suite_t *cipher_suite)
 {
-    if (operation->stage != PSA_PAKE_OPERATION_STAGE_COLLECT_INPUTS) {
-        return PSA_ERROR_BAD_STATE;
-    }
-
-    if (operation->alg != PSA_ALG_NONE) {
+    if (operation->stage != PSA_PAKE_OPERATION_STAGE_SETUP) {
         return PSA_ERROR_BAD_STATE;
     }
 
@@ -7266,6 +7262,8 @@ psa_status_t psa_pake_setup(
         computation_stage->output_step = PSA_PAKE_STEP_X1_X2;
     }
 
+    operation->stage = PSA_PAKE_OPERATION_STAGE_COLLECT_INPUTS;
+
     return PSA_SUCCESS;
 }
 
@@ -7281,10 +7279,6 @@ psa_status_t psa_pake_set_password_key(
         return PSA_ERROR_BAD_STATE;
     }
 
-    if (operation->alg == PSA_ALG_NONE) {
-        return PSA_ERROR_BAD_STATE;
-    }
-
     status = psa_get_and_lock_key_slot_with_policy(password, &slot,
                                                    PSA_KEY_USAGE_DERIVE,
                                                    PSA_ALG_JPAKE);
@@ -7329,10 +7323,6 @@ psa_status_t psa_pake_set_user(
         return PSA_ERROR_BAD_STATE;
     }
 
-    if (operation->alg == PSA_ALG_NONE) {
-        return PSA_ERROR_BAD_STATE;
-    }
-
     if (user_id_len == 0) {
         return PSA_ERROR_INVALID_ARGUMENT;
     }
@@ -7351,10 +7341,6 @@ psa_status_t psa_pake_set_peer(
         return PSA_ERROR_BAD_STATE;
     }
 
-    if (operation->alg == PSA_ALG_NONE) {
-        return PSA_ERROR_BAD_STATE;
-    }
-
     if (peer_id_len == 0) {
         return PSA_ERROR_INVALID_ARGUMENT;
     }
@@ -7370,10 +7356,6 @@ psa_status_t psa_pake_set_role(
         return PSA_ERROR_BAD_STATE;
     }
 
-    if (operation->alg == PSA_ALG_NONE) {
-        return PSA_ERROR_BAD_STATE;
-    }
-
     if (role != PSA_PAKE_ROLE_NONE &&
         role != PSA_PAKE_ROLE_FIRST &&
         role != PSA_PAKE_ROLE_SECOND &&
@@ -7887,7 +7869,7 @@ psa_status_t psa_pake_abort(
     }
 
     operation->alg = PSA_ALG_NONE;
-    operation->stage = PSA_PAKE_OPERATION_STAGE_COLLECT_INPUTS;
+    operation->stage = PSA_PAKE_OPERATION_STAGE_SETUP;
     operation->id = 0;
 
     return PSA_SUCCESS;