diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index f86ea3e..2039c1d 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -7767,8 +7767,8 @@
         psa_jpake_computation_stage_t *computation_stage =
             &operation->computation_stage.jpake;
 
-        computation_stage->round = FIRST;
-        computation_stage->mode = INPUT;
+        computation_stage->round = PSA_JPAKE_FIRST;
+        computation_stage->mode = PSA_JPAKE_INPUT;
         computation_stage->inputs = 0;
         computation_stage->outputs = 0;
         computation_stage->step = PSA_PAKE_STEP_KEY_SHARE;
@@ -7945,9 +7945,9 @@
 static psa_crypto_driver_pake_step_t convert_jpake_computation_stage_to_driver_step(
     psa_jpake_computation_stage_t *stage)
 {
-    if (stage->round == FIRST) {
+    if (stage->round == PSA_JPAKE_FIRST) {
         int is_x1;
-        if (stage->mode == OUTPUT) {
+        if (stage->mode == PSA_JPAKE_OUTPUT) {
             is_x1 = (stage->outputs < 1);
         } else {
             is_x1 = (stage->inputs < 1);
@@ -7976,8 +7976,8 @@
                     return PSA_JPAKE_STEP_INVALID;
             }
         }
-    } else if (stage->round == SECOND) {
-        if (stage->mode == OUTPUT) {
+    } else if (stage->round == PSA_JPAKE_SECOND) {
+        if (stage->mode == PSA_JPAKE_OUTPUT) {
             switch (stage->step) {
                 case PSA_PAKE_STEP_KEY_SHARE:
                     return PSA_JPAKE_X2S_STEP_KEY_SHARE;
@@ -8042,8 +8042,8 @@
             operation->stage = PSA_PAKE_OPERATION_STAGE_COMPUTATION;
             psa_jpake_computation_stage_t *computation_stage =
                 &operation->computation_stage.jpake;
-            computation_stage->round = FIRST;
-            computation_stage->mode = INPUT;
+            computation_stage->round = PSA_JPAKE_FIRST;
+            computation_stage->mode = PSA_JPAKE_INPUT;
             computation_stage->inputs = 0;
             computation_stage->outputs = 0;
             computation_stage->step = PSA_PAKE_STEP_KEY_SHARE;
@@ -8071,8 +8071,8 @@
     psa_jpake_computation_stage_t *computation_stage =
         &operation->computation_stage.jpake;
 
-    if (computation_stage->round != FIRST &&
-        computation_stage->round != SECOND) {
+    if (computation_stage->round != PSA_JPAKE_FIRST &&
+        computation_stage->round != PSA_JPAKE_SECOND) {
         return PSA_ERROR_BAD_STATE;
     }
 
@@ -8095,7 +8095,7 @@
 
     /* Check that we do not already have enough inputs/outputs
      * this round */
-    if (function_mode == INPUT) {
+    if (function_mode == PSA_JPAKE_INPUT) {
         if (computation_stage->inputs >=
             PSA_JPAKE_EXPECTED_INPUTS(computation_stage->round)) {
             return PSA_ERROR_BAD_STATE;
@@ -8118,16 +8118,16 @@
 
     if (stage->step == PSA_PAKE_STEP_ZK_PROOF) {
         /* End of an input/output */
-        if (function_mode == INPUT) {
+        if (function_mode == PSA_JPAKE_INPUT) {
             stage->inputs++;
             if (stage->inputs >= PSA_JPAKE_EXPECTED_INPUTS(stage->round)) {
-                stage->mode = OUTPUT;
+                stage->mode = PSA_JPAKE_OUTPUT;
             }
         }
-        if (function_mode == OUTPUT) {
+        if (function_mode == PSA_JPAKE_OUTPUT) {
             stage->outputs++;
             if (stage->outputs >= PSA_JPAKE_EXPECTED_OUTPUTS(stage->round)) {
-                stage->mode = INPUT;
+                stage->mode = PSA_JPAKE_INPUT;
             }
         }
         if (stage->inputs >= PSA_JPAKE_EXPECTED_INPUTS(stage->round) &&
@@ -8177,7 +8177,7 @@
     switch (operation->alg) {
 #if defined(PSA_WANT_ALG_JPAKE)
         case PSA_ALG_JPAKE:
-            status = psa_jpake_prologue(operation, step, OUTPUT);
+            status = psa_jpake_prologue(operation, step, PSA_JPAKE_OUTPUT);
             if (status != PSA_SUCCESS) {
                 goto exit;
             }
@@ -8201,7 +8201,7 @@
     switch (operation->alg) {
 #if defined(PSA_WANT_ALG_JPAKE)
         case PSA_ALG_JPAKE:
-            status = psa_jpake_epilogue(operation, OUTPUT);
+            status = psa_jpake_epilogue(operation, PSA_JPAKE_OUTPUT);
             if (status != PSA_SUCCESS) {
                 goto exit;
             }
@@ -8250,7 +8250,7 @@
     switch (operation->alg) {
 #if defined(PSA_WANT_ALG_JPAKE)
         case PSA_ALG_JPAKE:
-            status = psa_jpake_prologue(operation, step, INPUT);
+            status = psa_jpake_prologue(operation, step, PSA_JPAKE_INPUT);
             if (status != PSA_SUCCESS) {
                 goto exit;
             }
@@ -8274,7 +8274,7 @@
     switch (operation->alg) {
 #if defined(PSA_WANT_ALG_JPAKE)
         case PSA_ALG_JPAKE:
-            status = psa_jpake_epilogue(operation, INPUT);
+            status = psa_jpake_epilogue(operation, PSA_JPAKE_INPUT);
             if (status != PSA_SUCCESS) {
                 goto exit;
             }
@@ -8309,7 +8309,7 @@
     if (operation->alg == PSA_ALG_JPAKE) {
         psa_jpake_computation_stage_t *computation_stage =
             &operation->computation_stage.jpake;
-        if (computation_stage->round != FINISHED) {
+        if (computation_stage->round != PSA_JPAKE_FINISHED) {
             status = PSA_ERROR_BAD_STATE;
             goto exit;
         }