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 );