Fix ASRModelAdapter state passing to F() and parameter count

- Fixed state passing by extracting encoder_adapter and decoder_adapter from the full state dict in adapt_encoder_features and adapt_decoder_features
- Fixed expected parameter count from 33664 to 33600 in test_parameter_counts

Author: vishesh9131

axlearn/audio/adapter.py

@@ -218,7 +218,7 @@ def adapt_encoder_features(self, features, *, is_training=False, prng_key=None,
                 inputs=(features,),
                 is_training=is_training,
                 prng_key=prng_key,
-                state=state,
+                state=state["encoder_adapter"],
             )
             return outputs
 
@@ -248,7 +248,7 @@ def adapt_decoder_features(self, features, *, is_training=False, prng_key=None,
                 inputs=(features,),
                 is_training=is_training,
                 prng_key=prng_key,
-                state=state,
+                state=state["decoder_adapter"],
             )
             return outputs
 

axlearn/audio/adapter_test.py

@@ -191,7 +191,7 @@ def test_parameter_counts(self):
         total_params += np.prod(layer_norm_scale.shape)
         total_params += np.prod(layer_norm_bias.shape)
 
-        self.assertEqual(total_params, 33664)
+        self.assertEqual(total_params, 33600)
 
     @parameterized.parameters([True, False])
     def test_training_vs_eval_mode(self, is_training: bool):