Remove @nowrap in BaseEmbedding.init_states().

GitOrigin-RevId: ff83ff4

Author: ds-hwang

axlearn/common/decoder.py

@@ -686,7 +686,8 @@ def prefill_states(
         input_segment_ids = input_batch.get("input_segment_ids", None)
         positions = input_batch.get("positions", None)
 
-        emb = self.emb.init_states(batch_size=input_ids.shape[0], dtype=self.dtype())
+        with child_context("emb_init_states", module=self.emb):
+            emb = self.emb.init_states(batch_size=input_ids.shape[0], dtype=self.dtype())
         states, outputs = self._forward_for_mode(
             mode=ForwardMode.INIT_STATES,
             cached_states=dict(emb=emb, transformer_state=time_step),

axlearn/common/embedding.py

@@ -11,7 +11,7 @@
 from axlearn.common.base_layer import BaseLayer
 from axlearn.common.config import REQUIRED, InstantiableConfig, Required, config_class
 from axlearn.common.layers import Dropout, Embedding
-from axlearn.common.module import Module, Tensor, child_context, nowrap
+from axlearn.common.module import Module, Tensor, child_context
 from axlearn.common.utils import Nested, validate_contains_paths
 
 
@@ -50,7 +50,6 @@ def attend(self, x: Tensor):
         """
         raise NotImplementedError(type(self))
 
-    @nowrap
     def init_states(self, *, batch_size: int, dtype: jnp.dtype) -> Nested[Tensor]:
         """Initializes state for streaming decode.