diff --git a/library/psa_crypto_aead.c b/library/psa_crypto_aead.c
index 2632830..356679c 100644
--- a/library/psa_crypto_aead.c
+++ b/library/psa_crypto_aead.c
@@ -154,10 +154,14 @@
             return( PSA_ERROR_NOT_SUPPORTED );
     }
 
-    if( PSA_AEAD_TAG_LENGTH( alg ) > full_tag_length )
+    if( PSA_AEAD_TAG_LENGTH( attributes->core.type,
+                             key_bits, alg )
+        > full_tag_length )
         return( PSA_ERROR_INVALID_ARGUMENT );
 
-    operation->tag_length = PSA_AEAD_TAG_LENGTH( alg );
+    operation->tag_length = PSA_AEAD_TAG_LENGTH( attributes->core.type,
+                                                 key_bits,
+                                                 alg );
 
     return( PSA_SUCCESS );
 }