haliax 1.4.dev445__tar.gz → 1.4.dev446__tar.gz

{haliax-1.4.dev445 → haliax-1.4.dev446}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev445
+Version: 1.4.dev446
 Summary: Named Tensors for Legible Deep Learning in JAX
 Project-URL: Homepage, https://github.com/stanford-crfm/haliax
 Project-URL: Bug Tracker, https://github.com/stanford-crfm/haliax/issues/

{haliax-1.4.dev445 → haliax-1.4.dev446}/src/haliax/__about__.py

@@ -3,4 +3,4 @@
 # SPDX-License-Identifier: Apache-2.0
 
 
-__version__ = "1.4.dev445"
+__version__ = "1.4.dev446"

{haliax-1.4.dev445 → haliax-1.4.dev446}/src/haliax/nn/scan.py

@@ -150,6 +150,11 @@ class BlockFoldable(Protocol[M]):
         """
         ...
 
+    def get_layer(self, index: int) -> M:
+        """Return the ``index``th layer of the folded module."""
+
+        ...
+
 
 class BlockSeq(ModuleWithStateDictSerialization, Generic[M]):
     """
@@ -377,6 +382,11 @@ class BlockSeq(ModuleWithStateDictSerialization, Generic[M]):
 
         return state_dict
 
+    def get_layer(self, index: int) -> M:
+        """Return the ``index``th block in this sequential container."""
+
+        return self.blocks[index]
+
     @property
     def _output_ckpt_name(self):
         return f"BlockSeq[{self.Block}, {self.blocks[0].__class__.__name__}].outputs"
@@ -442,6 +452,12 @@ class Stacked(ModuleWithStateDictSerialization, Generic[M]):
     Block: Axis = eqx.field(static=True)
     gradient_checkpointing: ScanCheckpointPolicy = eqx.field(static=True)
 
+    @property
+    def Layers(self) -> Axis:
+        """Alias for :attr:`Block` used by some downstream code."""
+
+        return self.Block
+
     @classmethod
     def init(
         cls,
@@ -700,6 +716,23 @@ class Stacked(ModuleWithStateDictSerialization, Generic[M]):
         unstacked_leaves = tuple(zip(*unstacked_leaves))
         return tuple(map(lambda x: jax.tree_util.tree_unflatten(structure, x), unstacked_leaves))
 
+    def get_layer(self, index: int) -> M:
+        """Return the ``index``th layer of this stacked module."""
+
+        def select_leaf(leaf):
+            if isinstance(leaf, haliax.NamedArray):
+                if haliax.selects_axis(leaf.axes, self.Block):
+                    return leaf[self.Block, index]
+                else:
+                    return leaf
+            if is_jax_or_hax_array_like(leaf):
+                if getattr(leaf, "shape", ()) and leaf.shape[0] == self.Block.size:
+                    return leaf[index]
+                return leaf
+            return leaf
+
+        return haliax.tree_util.tree_map(select_leaf, self.stacked)
+
     def to_state_dict(self, prefix: str | None = None) -> StateDict:
         # this method needs to "devectorize" the blocks, so that we have a list of blocks h.0.FOO, h.1.FOO, etc.
         # first just do the normal thing with our own dict, which we'll post-process

{haliax-1.4.dev445 → haliax-1.4.dev446}/tests/test_scan.py

@@ -51,6 +51,65 @@ def test_unstacked():
         assert hax.all(module.array == m.stacked.array[i])
 
 
+def test_get_layer_stacked():
+    class Module(eqx.Module):
+        named: hax.NamedArray
+        array: jax.Array
+        static: int = eqx.static_field()
+
+        def __call__(self, x, *, key):  # pragma: no cover - unused in this test
+            return x + self.array + self.static
+
+        @staticmethod
+        def init(named, array, static):
+            return Module(named=named, array=array, static=static)
+
+    Block = hax.Axis("block", 3)
+    E = hax.Axis("E", 4)
+
+    initial_named = hax.random.uniform(jax.random.PRNGKey(0), (Block, E))
+    arrays = jax.numpy.arange(Block.size)
+
+    stacked = Stacked.init(Block, Module)(named=initial_named, array=arrays, static=2)
+
+    layer = stacked.get_layer(1)
+
+    assert isinstance(layer, Module)
+    assert layer.static == 2
+    assert layer.named.axes == (E,)
+    assert hax.all(layer.named == initial_named["block", 1])
+    assert hax.all(layer.array == arrays[1])
+
+
+def test_get_layer_blockseq():
+    class Module(eqx.Module):
+        named: hax.NamedArray
+        array: jax.Array
+        static: int = eqx.static_field()
+
+        def __call__(self, x, *, key):  # pragma: no cover - unused in this test
+            return x + self.array + self.static
+
+        @staticmethod
+        def init(named, array, static):
+            return Module(named=named, array=array, static=static)
+
+    Block = hax.Axis("block", 3)
+    E = hax.Axis("E", 4)
+
+    initial_named = hax.random.uniform(jax.random.PRNGKey(0), (Block, E))
+    arrays = jax.numpy.arange(Block.size)
+
+    seq = BlockSeq.init(Block, Module)(named=initial_named, array=arrays, static=2)
+
+    layer = seq.get_layer(2)
+
+    assert isinstance(layer, Module)
+    assert layer.static == 2
+    assert hax.all(layer.named == initial_named["block", 2])
+    assert hax.all(layer.array == arrays[2])
+
+
 def test_vmap():
     class Module(eqx.Module):
         weight: hax.NamedArray