diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index f3a2c64..ab9ec72 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -1225,7 +1225,7 @@
     mbedtls_cipher_id_t cipher_id_tmp;
 
     if( PSA_ALG_IS_AEAD( alg ) )
-        alg &= ~PSA_ALG_AEAD_TAG_LENGTH_MASK;
+        alg = PSA_ALG_AEAD_WITH_TAG_LENGTH( alg, 0 );
 
     if( PSA_ALG_IS_CIPHER( alg ) || PSA_ALG_IS_AEAD( alg ) )
     {
@@ -1249,10 +1249,10 @@
             case PSA_ALG_CBC_PKCS7:
                 mode = MBEDTLS_MODE_CBC;
                 break;
-            case PSA_ALG_CCM & ~PSA_ALG_AEAD_TAG_LENGTH_MASK:
+            case PSA_ALG_AEAD_WITH_TAG_LENGTH( PSA_ALG_CCM, 0 ):
                 mode = MBEDTLS_MODE_CCM;
                 break;
-            case PSA_ALG_GCM & ~PSA_ALG_AEAD_TAG_LENGTH_MASK:
+            case PSA_ALG_AEAD_WITH_TAG_LENGTH( PSA_ALG_GCM, 0 ):
                 mode = MBEDTLS_MODE_GCM;
                 break;
             default: