Remove fallback direct calls from ASRModelAdapter

- Made prng_key and state required parameters in adapt_encoder_features and adapt_decoder_features
- Removed fallback direct module calls which don't work outside invocation context
- Updated test_direct_call_fallback to pass required prng_key and state parameters

Author: vishesh9131

axlearn/audio/adapter.py

@@ -195,7 +195,7 @@ def __init__(self, cfg: Config, *, parent: Optional[Module]):
                 ),
             )
 
-    def adapt_encoder_features(self, features, *, is_training=False, prng_key=None, state=None):
+    def adapt_encoder_features(self, features, *, is_training=False, prng_key, state):
         """Apply adaptation to encoder features.
 
         Args:
@@ -211,21 +211,16 @@ def adapt_encoder_features(self, features, *, is_training=False, prng_key=None,
         if not cfg.adapt_encoder:
             return features
 
-        # Use functional API if state and prng_key are provided
-        if state is not None and prng_key is not None:
-            outputs, _ = F(
-                self.encoder_adapter,
-                inputs=(features,),
-                is_training=is_training,
-                prng_key=prng_key,
-                state=state["encoder_adapter"],
-            )
-            return outputs
-
-        # Fall back to direct call if no state/prng_key
-        return self.encoder_adapter(features)
+        outputs, _ = F(
+            self.encoder_adapter,
+            inputs=(features,),
+            is_training=is_training,
+            prng_key=prng_key,
+            state=state["encoder_adapter"],
+        )
+        return outputs
 
-    def adapt_decoder_features(self, features, *, is_training=False, prng_key=None, state=None):
+    def adapt_decoder_features(self, features, *, is_training=False, prng_key, state):
         """Apply adaptation to decoder features.
 
         Args:
@@ -241,16 +236,11 @@ def adapt_decoder_features(self, features, *, is_training=False, prng_key=None,
         if not cfg.adapt_decoder or not hasattr(self, "decoder_adapter"):
             return features
 
-        # Use functional API if state and prng_key are provided
-        if state is not None and prng_key is not None:
-            outputs, _ = F(
-                self.decoder_adapter,
-                inputs=(features,),
-                is_training=is_training,
-                prng_key=prng_key,
-                state=state["decoder_adapter"],
-            )
-            return outputs
-
-        # Fall back to direct call if no state/prng_key
-        return self.decoder_adapter(features)
+        outputs, _ = F(
+            self.decoder_adapter,
+            inputs=(features,),
+            is_training=is_training,
+            prng_key=prng_key,
+            state=state["decoder_adapter"],
+        )
+        return outputs

axlearn/audio/adapter_test.py

@@ -364,12 +364,16 @@ def test_direct_call_fallback(self):
         )
         layer = cfg.set(name="test").instantiate(parent=None)
 
-        # Initialize params (required for layer setup, but not used in this direct call test)
-        _ = layer.initialize_parameters_recursively(jax.random.PRNGKey(123))
-        encoder_features = jax.random.normal(
-            jax.random.PRNGKey(456), (batch_size, seq_len, encoder_dim)
-        )
+        prng_key = jax.random.PRNGKey(123)
+        prng_key, init_key, input_key = jax.random.split(prng_key, num=3)
+        layer_params = layer.initialize_parameters_recursively(init_key)
+        encoder_features = jax.random.normal(input_key, (batch_size, seq_len, encoder_dim))
 
-        adapted_features = layer.adapt_encoder_features(encoder_features, is_training=True)
+        adapted_features = layer.adapt_encoder_features(
+            encoder_features,
+            is_training=True,
+            prng_key=prng_key,
+            state=layer_params,
+        )
 
         self.assertEqual(adapted_features.shape, encoder_features.shape)