diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 83b45f097..65dc5c7fe 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -3901,6 +3901,41 @@ psa_status_t psa_aead_set_lengths( psa_aead_operation_t *operation,
         goto exit;
     }
 
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_GCM)
+    if( operation->alg == PSA_ALG_GCM )
+    {
+        /* Lengths can only be too large for GCM if size_t is bigger than 32
+        * bits. Without the guard this code will generate warnings on 32bit
+        * builds */
+#if SIZE_MAX > UINT32_MAX
+        if( (( uint64_t ) ad_length ) >> 61 != 0 ||
+            (( uint64_t ) plaintext_length ) > 0xFFFFFFFE0ull )
+        {
+            status = PSA_ERROR_INVALID_ARGUMENT;
+            goto exit;
+        }
+#endif
+    }
+    else
+#endif /* MBEDTLS_PSA_BUILTIN_ALG_GCM */
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_CCM)
+    if( operation->alg == PSA_ALG_CCM )
+    {
+        if( ad_length > 0xFF00 )
+        {
+            status = PSA_ERROR_INVALID_ARGUMENT;
+            goto exit;
+        }
+    }
+    else
+#endif /* MBEDTLS_PSA_BUILTIN_ALG_CCM */
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_CHACHA20_POLY1305)
+    if( operation->alg == PSA_ALG_CHACHA20_POLY1305 )
+    {
+        /* No length restrictions for ChaChaPoly. */
+    }
+#endif /* MBEDTLS_PSA_BUILTIN_ALG_CHACHA20_POLY1305 */
+
     status = psa_driver_wrapper_aead_set_lengths( operation, ad_length,
                                                   plaintext_length );
 
diff --git a/library/psa_crypto_aead.c b/library/psa_crypto_aead.c
index 9000abf30..d7317bd1a 100644
--- a/library/psa_crypto_aead.c
+++ b/library/psa_crypto_aead.c
@@ -477,55 +477,6 @@ psa_status_t mbedtls_psa_aead_set_nonce(
     return( status );
 }
 
- /* Declare the lengths of the message and additional data for AEAD. */
-psa_status_t mbedtls_psa_aead_set_lengths(
-    mbedtls_psa_aead_operation_t *operation,
-    size_t ad_length,
-    size_t plaintext_length )
-{
-
-#if defined(MBEDTLS_PSA_BUILTIN_ALG_GCM)
-    if( operation->alg == PSA_ALG_GCM )
-    {
-        /* Lengths can only be too large for GCM if size_t is bigger than 32
-         * bits. Without the guard this code will generate warnings on 32bit
-         * builds */
-#if SIZE_MAX > UINT32_MAX
-        if( ( (uint64_t) ad_length ) >> 61 != 0 ||
-            ( (uint64_t) plaintext_length ) > 0xFFFFFFFE0ull )
-        {
-            return ( PSA_ERROR_INVALID_ARGUMENT );
-        }
-#endif
-    }
-    else
-#endif /* MBEDTLS_PSA_BUILTIN_ALG_GCM */
-#if defined(MBEDTLS_PSA_BUILTIN_ALG_CCM)
-    if( operation->alg == PSA_ALG_CCM )
-    {
-        if( ad_length > 0xFF00 )
-            return ( PSA_ERROR_INVALID_ARGUMENT );
-    }
-    else
-#endif /* MBEDTLS_PSA_BUILTIN_ALG_CCM */
-#if defined(MBEDTLS_PSA_BUILTIN_ALG_CHACHA20_POLY1305)
-    if( operation->alg == PSA_ALG_CHACHA20_POLY1305 )
-    {
-        /* No length restrictions for ChaChaPoly. */
-    }
-    else
-#endif /* MBEDTLS_PSA_BUILTIN_ALG_CHACHA20_POLY1305 */
-    {
-        ( void ) operation;
-        ( void ) ad_length;
-        ( void ) plaintext_length;
-
-        return ( PSA_ERROR_NOT_SUPPORTED );
-    }
-
-    return ( PSA_SUCCESS );
-}
-
 /* Pass additional data to an active multipart AEAD operation. */
 psa_status_t mbedtls_psa_aead_update_ad(
     mbedtls_psa_aead_operation_t *operation,
diff --git a/library/psa_crypto_aead.h b/library/psa_crypto_aead.h
index e82e1cc09..f968c15c8 100644
--- a/library/psa_crypto_aead.h
+++ b/library/psa_crypto_aead.h
@@ -267,47 +267,6 @@ psa_status_t mbedtls_psa_aead_set_nonce(
     const uint8_t *nonce,
     size_t nonce_length );
 
-/** Declare the lengths of the message and additional data for AEAD.
- *
- * \note The signature of this function is that of a PSA driver aead_set_lengths
- *       entry point. This function behaves as an aead_set_lengths entry point
- *       as defined in the PSA driver interface specification for transparent
- *       drivers.
- *
- * The PSA core calls this function before calling mbedtls_psa_aead_update_ad()
- * or mbedtls_psa_aead_update() if the algorithm for the operation requires it.
- * If the algorithm does not require it, calling this function is optional, but
- * if this function is called then the implementation must enforce the lengths.
- *
- * The PSA core may call this function before or after setting the nonce with
- * mbedtls_psa_aead_set_nonce().
- *
- * - For #PSA_ALG_CCM, calling this function is required.
- * - For the other AEAD algorithms defined in this specification, calling
- *   this function is not required.
- *
- * If this function returns an error status, the PSA core calls
- * mbedtls_psa_aead_abort().
- *
- * \param[in,out] operation     Active AEAD operation.
- * \param ad_length             Size of the non-encrypted additional
- *                              authenticated data in bytes.
- * \param plaintext_length      Size of the plaintext to encrypt in bytes.
- *
- * \retval #PSA_SUCCESS
- *         Success.
- * \retval #PSA_ERROR_INVALID_ARGUMENT
- *         At least one of the lengths is not acceptable for the chosen
- *         algorithm.
- * \retval #PSA_ERROR_NOT_SUPPORTED
- *         Algorithm previously set is not supported in this configuration of
- *         the library.
- */
-psa_status_t mbedtls_psa_aead_set_lengths(
-    mbedtls_psa_aead_operation_t *operation,
-    size_t ad_length,
-    size_t plaintext_length );
-
 /** Pass additional data to an active AEAD operation.
  *
  *  \note The signature of this function is that of a PSA driver
diff --git a/library/psa_crypto_driver_wrappers.c b/library/psa_crypto_driver_wrappers.c
index cfc77fbb5..4bbb61c3d 100644
--- a/library/psa_crypto_driver_wrappers.c
+++ b/library/psa_crypto_driver_wrappers.c
@@ -1706,9 +1706,9 @@ psa_status_t psa_driver_wrapper_aead_set_lengths(
     {
 #if defined(MBEDTLS_PSA_BUILTIN_AEAD)
         case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
-            return( mbedtls_psa_aead_set_lengths( &operation->ctx.mbedtls_ctx,
-                                                  ad_length,
-                                                  plaintext_length ) );
+            /* No mbedtls_psa_aead_set_lengths, everything is done in PSA
+             * Core. */
+            return( PSA_SUCCESS );
 
 #endif /* MBEDTLS_PSA_BUILTIN_AEAD */
 
diff --git a/tests/src/drivers/test_driver_aead.c b/tests/src/drivers/test_driver_aead.c
index 6befe7cc0..d27ada294 100644
--- a/tests/src/drivers/test_driver_aead.c
+++ b/tests/src/drivers/test_driver_aead.c
@@ -171,9 +171,8 @@ psa_status_t mbedtls_test_transparent_aead_set_lengths(
     }
     else
     {
-        mbedtls_test_driver_aead_hooks.driver_status =
-            mbedtls_psa_aead_set_lengths( operation, ad_length,
-                                          plaintext_length );
+       /* No mbedtls_psa_aead_set_lengths, everything is done in PSA Core. */
+       mbedtls_test_driver_aead_hooks.driver_status = PSA_SUCCESS;
     }
 
     return( mbedtls_test_driver_aead_hooks.driver_status );