diff --git a/include/psa/crypto.h b/include/psa/crypto.h
index 2046947..2c3288e 100644
--- a/include/psa/crypto.h
+++ b/include/psa/crypto.h
@@ -147,6 +147,7 @@
  * by the following functions:
  * - psa_make_key_persistent()
  * - psa_set_key_type()
+ * - psa_set_key_bits()
  * - psa_set_key_usage_flags()
  * - psa_set_key_algorithm()
  * - psa_reset_key_attributes()
@@ -293,6 +294,20 @@
 static void psa_set_key_type(psa_key_attributes_t *attributes,
                              psa_key_type_t type);
 
+/** Declare the size of a key.
+ *
+ * This function overwrites any key size previously set in \p attributes.
+ *
+ * This function may be declared as `static` (i.e. without external
+ * linkage). This function may be provided as a function-like macro,
+ * but in this case it must evaluate each of its arguments exactly once.
+ *
+ * \param[out] attributes       The attribute structure to write to.
+ * \param bits                  The key size in bits.
+ */
+static void psa_set_key_bits(psa_key_attributes_t *attributes,
+                             size_t bits);
+
 /** Retrieve the key type from key attributes.
  *
  * This function may be declared as `static` (i.e. without external
@@ -331,11 +346,6 @@
  * property may not hold in future versions of this specification or
  * for implementation-specific values.
  *
- * In addition to the attributes that were set when creating the key,
- * this function reports the following data:
- * - The key size in bits, which can be retrieved with
- *   psa_get_key_bits().
- *
  * \param[in] handle            Handle to the key to query.
  * \param[in,out] attributes    On success, the attributes of the key.
  *                              On failure, equivalent to a
@@ -3018,12 +3028,8 @@
  * The generator's capacity is decreased by the number of bytes read.
  *
  * \param[in] attributes    The attributes for the new key.
- *                          The key size field in \p attributes is
- *                          ignored; the actual key size is taken
- *                          from the \p bits parameter instead.
  * \param[out] handle       On success, a handle to the newly created key.
  *                          \c 0 on failure.
- * \param bits              Key size in bits.
  * \param[in,out] generator The generator object to read from.
  *
  * \retval #PSA_SUCCESS
@@ -3054,7 +3060,6 @@
  */
 psa_status_t psa_generator_import_key(const psa_key_attributes_t *attributes,
                                       psa_key_handle_t *handle,
-                                      size_t bits,
                                       psa_crypto_generator_t *generator);
 
 /** Abort a generator.
@@ -3383,12 +3388,8 @@
  * \brief Generate a key or key pair.
  *
  * \param[in] attributes    The attributes for the new key.
- *                          The key size field in \p attributes is
- *                          ignored; the actual key size is taken
- *                          from the \p bits parameter instead.
  * \param[out] handle       On success, a handle to the newly created key.
  *                          \c 0 on failure.
- * \param bits              Key size in bits.
  * \param[in] extra         Extra parameters for key generation. The
  *                          interpretation of this parameter depends on
  *                          the key type \c type. All types support \c NULL to
@@ -3447,7 +3448,6 @@
  */
 psa_status_t psa_generate_key(const psa_key_attributes_t *attributes,
                               psa_key_handle_t *handle,
-                              size_t bits,
                               const void *extra,
                               size_t extra_size);
 
diff --git a/include/psa/crypto_struct.h b/include/psa/crypto_struct.h
index 273f6b6..f89073b 100644
--- a/include/psa/crypto_struct.h
+++ b/include/psa/crypto_struct.h
@@ -333,6 +333,12 @@
     return( attributes->type );
 }
 
+static inline void psa_set_key_bits(psa_key_attributes_t *attributes,
+                                    size_t bits)
+{
+    attributes->bits = bits;
+}
+
 static inline size_t psa_get_key_bits(
     const psa_key_attributes_t *attributes)
 {
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index e035ece..c1e3a3f 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -4222,7 +4222,6 @@
 
 psa_status_t psa_generator_import_key( const psa_key_attributes_t *attributes,
                                        psa_key_handle_t *handle,
-                                       size_t bits,
                                        psa_crypto_generator_t *generator )
 {
     psa_status_t status;
@@ -4230,7 +4229,9 @@
     status = psa_start_key_creation( attributes, handle, &slot );
     if( status == PSA_SUCCESS )
     {
-        status = psa_generator_import_key_internal( slot, bits, generator );
+        status = psa_generator_import_key_internal( slot,
+                                                    attributes->bits,
+                                                    generator );
     }
     if( status == PSA_SUCCESS )
         status = psa_finish_key_creation( slot );
@@ -5139,7 +5140,6 @@
 
 psa_status_t psa_generate_key( const psa_key_attributes_t *attributes,
                                psa_key_handle_t *handle,
-                               size_t bits,
                                const void *extra,
                                size_t extra_size )
 {
@@ -5148,7 +5148,8 @@
     status = psa_start_key_creation( attributes, handle, &slot );
     if( status == PSA_SUCCESS )
     {
-        status = psa_generate_key_internal( slot, bits, extra, extra_size );
+        status = psa_generate_key_internal( slot, attributes->bits,
+                                            extra, extra_size );
     }
     if( status == PSA_SUCCESS )
         status = psa_finish_key_creation( slot );
diff --git a/programs/psa/crypto_examples.c b/programs/psa/crypto_examples.c
index 07d1fd2..72fa12f 100644
--- a/programs/psa/crypto_examples.c
+++ b/programs/psa/crypto_examples.c
@@ -162,9 +162,9 @@
                              PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT );
     psa_set_key_algorithm( &attributes, alg );
     psa_set_key_type( &attributes, PSA_KEY_TYPE_AES );
+    psa_set_key_bits( &attributes, key_bits );
 
-    status = psa_generate_key( &attributes, &key_handle, key_bits,
-                               NULL, 0 );
+    status = psa_generate_key( &attributes, &key_handle, NULL, 0 );
     ASSERT_STATUS( status, PSA_SUCCESS );
 
     status = cipher_encrypt( key_handle, alg, iv, sizeof( iv ),
@@ -213,9 +213,9 @@
                              PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT );
     psa_set_key_algorithm( &attributes, alg );
     psa_set_key_type( &attributes, PSA_KEY_TYPE_AES );
+    psa_set_key_bits( &attributes, key_bits );
 
-    status = psa_generate_key( &attributes, &key_handle, key_bits,
-                               NULL, 0 );
+    status = psa_generate_key( &attributes, &key_handle, NULL, 0 );
     ASSERT_STATUS( status, PSA_SUCCESS );
 
     status = cipher_encrypt( key_handle, alg, iv, sizeof( iv ),
@@ -260,9 +260,9 @@
                              PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT );
     psa_set_key_algorithm( &attributes, alg );
     psa_set_key_type( &attributes, PSA_KEY_TYPE_AES );
+    psa_set_key_bits( &attributes, key_bits );
 
-    status = psa_generate_key( &attributes, &key_handle, key_bits,
-                               NULL, 0 );
+    status = psa_generate_key( &attributes, &key_handle, NULL, 0 );
     ASSERT_STATUS( status, PSA_SUCCESS );
 
     status = cipher_encrypt( key_handle, alg, iv, sizeof( iv ),
diff --git a/programs/psa/key_ladder_demo.c b/programs/psa/key_ladder_demo.c
index b84e7fd..c1e296f 100644
--- a/programs/psa/key_ladder_demo.c
+++ b/programs/psa/key_ladder_demo.c
@@ -206,10 +206,9 @@
                              PSA_KEY_USAGE_DERIVE | PSA_KEY_USAGE_EXPORT );
     psa_set_key_algorithm( &attributes, KDF_ALG );
     psa_set_key_type( &attributes, PSA_KEY_TYPE_DERIVE );
+    psa_set_key_bits( &attributes, PSA_BYTES_TO_BITS( KEY_SIZE_BYTES ) );
 
-    PSA_CHECK( psa_generate_key( &attributes, &key_handle,
-                                 PSA_BYTES_TO_BITS( KEY_SIZE_BYTES ),
-                                 NULL, 0 ) );
+    PSA_CHECK( psa_generate_key( &attributes, &key_handle, NULL, 0 ) );
 
     PSA_CHECK( save_key( key_handle, key_file_name ) );
 
@@ -287,6 +286,7 @@
                              PSA_KEY_USAGE_DERIVE | PSA_KEY_USAGE_EXPORT );
     psa_set_key_algorithm( &attributes, KDF_ALG );
     psa_set_key_type( &attributes, PSA_KEY_TYPE_DERIVE );
+    psa_set_key_bits( &attributes, PSA_BYTES_TO_BITS( KEY_SIZE_BYTES ) );
 
     /* For each label in turn, ... */
     for( i = 0; i < ladder_depth; i++ )
@@ -306,10 +306,8 @@
         *key_handle = 0;
         /* Use the generator obtained from the parent key to create
          * the next intermediate key. */
-        PSA_CHECK( psa_generator_import_key(
-                       &attributes, key_handle,
-                       PSA_BYTES_TO_BITS( KEY_SIZE_BYTES ),
-                       &generator ) );
+        PSA_CHECK( psa_generator_import_key( &attributes, key_handle,
+                                             &generator ) );
         PSA_CHECK( psa_generator_abort( &generator ) );
     }
 
@@ -336,6 +334,7 @@
     psa_set_key_usage_flags( &attributes, usage );
     psa_set_key_algorithm( &attributes, WRAPPING_ALG );
     psa_set_key_type( &attributes, PSA_KEY_TYPE_AES );
+    psa_set_key_bits( &attributes, WRAPPING_KEY_BITS );
 
     PSA_CHECK( psa_key_derivation(
                    &generator,
@@ -345,8 +344,7 @@
                    NULL, 0,
                    PSA_BITS_TO_BYTES( WRAPPING_KEY_BITS ) ) );
     PSA_CHECK( psa_generator_import_key( &attributes, wrapping_key_handle,
-                   WRAPPING_KEY_BITS,
-                   &generator ) );
+                                         &generator ) );
 
 exit:
     psa_generator_abort( &generator );
diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index 5b99b84..c91094c 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -2,7 +2,7 @@
 static_checks:
 
 PSA key attributes structure
-attributes_set_get:0x6963:PSA_KEY_LIFETIME_PERSISTENT:PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_CCM:PSA_KEY_TYPE_AES
+attributes_set_get:0x6963:PSA_KEY_LIFETIME_PERSISTENT:PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_CCM:PSA_KEY_TYPE_AES:128
 
 PSA import/export raw: 0 bytes
 import_export:"":PSA_KEY_TYPE_RAW_DATA:0:PSA_KEY_USAGE_EXPORT:0:0:PSA_SUCCESS:1
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 1d67c6d..152f7e9 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -1157,7 +1157,7 @@
 /* BEGIN_CASE */
 void attributes_set_get( int id_arg, int lifetime_arg,
                          int usage_flags_arg, int alg_arg,
-                         int type_arg )
+                         int type_arg, int bits_arg )
 {
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     psa_key_id_t id = id_arg;
@@ -1165,23 +1165,27 @@
     psa_key_usage_t usage_flags = usage_flags_arg;
     psa_algorithm_t alg = alg_arg;
     psa_key_type_t type = type_arg;
+    size_t bits = bits_arg;
 
     TEST_EQUAL( psa_get_key_id( &attributes ), 0 );
     TEST_EQUAL( psa_get_key_lifetime( &attributes ), 0 );
     TEST_EQUAL( psa_get_key_usage_flags( &attributes ), 0 );
     TEST_EQUAL( psa_get_key_algorithm( &attributes ), 0 );
     TEST_EQUAL( psa_get_key_type( &attributes ), 0 );
+    TEST_EQUAL( psa_get_key_bits( &attributes ), 0 );
 
     psa_make_key_persistent( &attributes, id, lifetime );
     psa_set_key_usage_flags( &attributes, usage_flags );
     psa_set_key_algorithm( &attributes, alg );
     psa_set_key_type( &attributes, type );
+    psa_set_key_bits( &attributes, bits );
 
     TEST_EQUAL( psa_get_key_id( &attributes ), id );
     TEST_EQUAL( psa_get_key_lifetime( &attributes ), lifetime );
     TEST_EQUAL( psa_get_key_usage_flags( &attributes ), usage_flags );
     TEST_EQUAL( psa_get_key_algorithm( &attributes ), alg );
     TEST_EQUAL( psa_get_key_type( &attributes ), type );
+    TEST_EQUAL( psa_get_key_bits( &attributes ), bits );
 
     psa_reset_key_attributes( &attributes );
 
@@ -1190,6 +1194,7 @@
     TEST_EQUAL( psa_get_key_usage_flags( &attributes ), 0 );
     TEST_EQUAL( psa_get_key_algorithm( &attributes ), 0 );
     TEST_EQUAL( psa_get_key_type( &attributes ), 0 );
+    TEST_EQUAL( psa_get_key_bits( &attributes ), 0 );
 }
 /* END_CASE */
 
@@ -4294,8 +4299,8 @@
     psa_set_key_usage_flags( &attributes, derived_usage );
     psa_set_key_algorithm( &attributes, derived_alg );
     psa_set_key_type( &attributes, derived_type );
+    psa_set_key_bits( &attributes, derived_bits );
     PSA_ASSERT( psa_generator_import_key( &attributes, &derived_handle,
-                                          derived_bits,
                                           &generator ) );
 
     /* Test the key information */
@@ -4327,7 +4332,6 @@
     psa_key_handle_t derived_handle = 0;
     psa_algorithm_t alg = alg_arg;
     size_t bytes1 = bytes1_arg;
-    size_t derived_bits = PSA_BYTES_TO_BITS( bytes1 );
     size_t bytes2 = bytes2_arg;
     size_t capacity = bytes1 + bytes2;
     psa_crypto_generator_t generator = PSA_CRYPTO_GENERATOR_INIT;
@@ -4365,16 +4369,16 @@
     psa_set_key_usage_flags( &derived_attributes, PSA_KEY_USAGE_EXPORT );
     psa_set_key_algorithm( &derived_attributes, 0 );
     psa_set_key_type( &derived_attributes, PSA_KEY_TYPE_RAW_DATA );
+    psa_set_key_bits( &derived_attributes, PSA_BYTES_TO_BITS( bytes1 ) );
     PSA_ASSERT( psa_generator_import_key( &derived_attributes, &derived_handle,
-                                          derived_bits,
                                           &generator ) );
     PSA_ASSERT( psa_export_key( derived_handle,
                                 export_buffer, bytes1,
                                 &length ) );
     TEST_EQUAL( length, bytes1 );
     PSA_ASSERT( psa_destroy_key( derived_handle ) );
+    psa_set_key_bits( &derived_attributes, PSA_BYTES_TO_BITS( bytes2 ) );
     PSA_ASSERT( psa_generator_import_key( &derived_attributes, &derived_handle,
-                                          PSA_BYTES_TO_BITS( bytes2 ),
                                           &generator ) );
     PSA_ASSERT( psa_export_key( derived_handle,
                                 export_buffer + bytes1, bytes2,
@@ -4667,9 +4671,10 @@
     psa_set_key_usage_flags( &attributes, usage );
     psa_set_key_algorithm( &attributes, alg );
     psa_set_key_type( &attributes, type );
+    psa_set_key_bits( &attributes, bits );
 
     /* Generate a key */
-    TEST_EQUAL( psa_generate_key( &attributes, &handle, bits, NULL, 0 ),
+    TEST_EQUAL( psa_generate_key( &attributes, &handle, NULL, 0 ),
                 expected_status );
     if( expected_info_status != PSA_SUCCESS )
         goto exit;
@@ -4722,6 +4727,7 @@
     psa_set_key_usage_flags( &attributes, usage_flags );
     psa_set_key_algorithm( &attributes, alg );
     psa_set_key_type( &attributes, type );
+    psa_set_key_bits( &attributes, bits );
 
     switch( generation_method )
     {
@@ -4733,8 +4739,7 @@
 
         case GENERATE_KEY:
             /* Generate a key */
-            PSA_ASSERT( psa_generate_key( &attributes, &handle,
-                                          bits, NULL, 0 ) );
+            PSA_ASSERT( psa_generate_key( &attributes, &handle, NULL, 0 ) );
             break;
 
         case DERIVE_KEY:
@@ -4757,7 +4762,7 @@
                                 &generator, PSA_KDF_STEP_INFO,
                                 NULL, 0 ) );
                 PSA_ASSERT( psa_generator_import_key( &attributes, &handle,
-                                                      bits, &generator ) );
+                                                      &generator ) );
                 PSA_ASSERT( psa_generator_abort( &generator ) );
                 PSA_ASSERT( psa_destroy_key( base_key ) );
                 base_key = 0;