haliax 1.4.dev308__tar.gz → 1.4.dev313__tar.gz

{haliax-1.4.dev308 → haliax-1.4.dev313}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: haliax
-Version: 1.4.dev308
+Version: 1.4.dev313
 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.dev313/src/haliax/__about__.py

@@ -0,0 +1 @@
+__version__ = "1.4.dev313"

{haliax-1.4.dev308 → haliax-1.4.dev313}/src/haliax/core.py

@@ -684,6 +684,8 @@ def take(array: NamedArray, axis: AxisSelector, index: Union[int, NamedArray]) -
         new_axes = array.axes[:axis_index] + array.axes[axis_index + 1 :]
         return NamedArray(new_array, new_axes)
     else:
+        if not jnp.issubdtype(index.dtype, jnp.integer):
+            raise ValueError(f"Index must be an integer array, got {index.dtype}")
         # #13: should broadcast/autobatch take
         remaining_axes = eliminate_axes(array.axes, axis)
         # axis order is generally [array.axes[:axis_index], index.axes, array.axes[axis_index + 1 :]]

{haliax-1.4.dev308 → haliax-1.4.dev313}/src/haliax/nn/scan.py

@@ -156,6 +156,13 @@ class Stacked(eqx.Module, Generic[M]):
 
     Stacked also supports gradient checkpointing, which is useful for very large models that don't fit in memory.
 
+    Typically only one of "fold" or "scan" can be used with a given Stacked module, depending on the what the module
+    returns: if the module returns a single output, use "fold"; if the module returns a sequence of intermediates and
+    an output to be passed to the next layer, use "scan". More concretely, for a transformer, you would use "scan" if
+    you wanted to return a kv cache (or the attention matrix) as well as the output of the transformer. If you just
+    wanted the output of the transformer, you would use "fold".
+
+
     Example:
         ```python
         >>> import equinox as eqx
@@ -206,6 +213,32 @@ class Stacked(eqx.Module, Generic[M]):
         return fn
 
     def scan(self, init, *extra_args, **extra_kwargs):
+        """
+        Scan over the stacked module. This is the same as a for loop that applies each instance of the module in sequence
+        to the input, passing the output of one instance to the next instance. It returns a stack of intermediates as
+        well as the final output.
+
+        That is, it behaves similarly to the following Python code:
+
+        ```python
+        carry = init
+        intermediates = []
+
+        for block in self.stacked:
+            carry, extra = block(carry)
+            intermediates.append(extra)
+
+        return carry, hax.stack(Block, intermediates)
+        ```
+
+        Args:
+            init:
+            *extra_args:
+            **extra_kwargs:
+
+        Returns:
+
+        """
         if self.gradient_checkpointing:
             do_block = filter_checkpoint(self._do_block, prevent_cse=self.prevent_cse)
         else:
@@ -213,6 +246,27 @@ class Stacked(eqx.Module, Generic[M]):
         return haliax.scan(do_block, self.Block)(init, self.stacked, *extra_args, **extra_kwargs)
 
     def fold(self, init, *args, **kwargs):
+        """
+        Fold over the stacked module. This is the same as a for loop that applies each instance of the module in sequence
+        to the input, passing the output of one instance to the next instance.
+        That is, it behaves similarly to the following Python code:
+
+        ```python
+        carry = init
+        for block in self.stacked:
+            carry = block(carry)
+
+        return carry
+        ```
+
+        Args:
+            init:
+            *args:
+            **kwargs:
+
+        Returns:
+
+        """
         if self.gradient_checkpointing:
             do_block = filter_checkpoint(self._do_block, prevent_cse=self.prevent_cse)
         else:

{haliax-1.4.dev308 → haliax-1.4.dev313}/src/haliax/partitioning.py

@@ -603,7 +603,10 @@ def round_axis_for_partitioning(axis: Axis, mapping: Optional[ResourceMapping] =
 
 
 def _get_mesh() -> Mesh:
-    from jax.experimental.maps import thread_resources
+    try:
+        from jax.interpreters.pxla import thread_resources
+    except ImportError:
+        from jax.experimental.maps import thread_resources
 
     return thread_resources.env.physical_mesh
 

{haliax-1.4.dev308 → haliax-1.4.dev313}/src/haliax/tree_util.py

@@ -1,3 +1,5 @@
+import dataclasses
+import functools
 from typing import Optional
 
 import equinox as eqx
@@ -5,6 +7,8 @@ import jax
 import jax.tree_util as jtu
 from jaxtyping import PRNGKeyArray, PyTree
 
+import haliax.nn
+
 from .axis import AxisSelector
 from .core import NamedArray
 from .jax_utils import maybe_rng_split
@@ -15,12 +19,45 @@ def tree_map(fn, tree, *rest, is_leaf=None):
     """
     Version of [jax.tree_util.tree_map][] that automatically treats NamedArrays as leaves.
     """
+    old_is_leaf = is_leaf
     if is_leaf is None:
         is_leaf = lambda x: isinstance(x, NamedArray)
     else:
-        is_leaf = lambda x: is_leaf(x) or is_named_array(x)
+        is_leaf = lambda x: old_is_leaf(x) or is_named_array(x)
+
+    return jax.tree.map(fn, tree, *rest, is_leaf=is_leaf)
+
+
+def scan_aware_tree_map(fn, tree, *rest, is_leaf=None):
+    """
+    Version of [haliax.tree_util.tree_map][] that is aware of the scan-layer pattern, specifically as implmeneted
+    in hax.nn.Stacked. This function will (implicitly) apply the transform to each layer in each Stacked module
+    (using vmap). If there are no Stacked modules in the tree, this function is equivalent to [haliax.tree_util.tree_map][].
+
+    Args:
+        fn:
+        tree:
+        *rest:
+        is_leaf:
+
+    Returns:
+    """
+    old_is_leaf = is_leaf
+    if is_leaf is None:
+        is_leaf = lambda x: isinstance(x, haliax.nn.Stacked)
+    else:
+        is_leaf = lambda x: old_is_leaf(x) or isinstance(x, haliax.nn.Stacked)
+
+    mapped_fn = functools.partial(scan_aware_tree_map, fn, is_leaf=is_leaf)
+
+    def rec_fn(x, *rest):
+        if isinstance(x, haliax.nn.Stacked):
+            new_inner = haliax.vmap(mapped_fn, x.Block)(x.stacked, *[r.stacked for r in rest])
+            return dataclasses.replace(x, stacked=new_inner)  # type: ignore
+        else:
+            return fn(x)
 
-    return jax.tree_util.tree_map(fn, tree, *rest, is_leaf=is_leaf)
+    return tree_map(rec_fn, tree, *rest, is_leaf=is_leaf)
 
 
 def tree_flatten(tree, is_leaf=None):

{haliax-1.4.dev308 → haliax-1.4.dev313}/tests/test_scan.py

@@ -68,3 +68,30 @@ def test_seq_and_stacked_give_same_results():
     y_seq = m_seq.fold(x, key=jax.random.split(jax.random.PRNGKey(2), Block.size))
 
     assert hax.all(hax.isclose(y, y_seq, atol=1e-5))
+
+
+def test_using_scan():
+    class Module(eqx.Module):
+        named: hax.NamedArray
+        array: jax.Array
+        static: int = eqx.static_field()
+
+        def __call__(self, x, *, key):
+            return x + self.array + self.static + hax.random.normal(key, x.axes), x * 2
+
+        @staticmethod
+        def init(named, array, static):
+            return Module(named=named, array=array, static=static)
+
+    Block = hax.Axis("block", 4)
+    E = hax.Axis("E", 10)
+
+    initial_named = hax.random.uniform(jax.random.PRNGKey(0), (Block, E))
+
+    m = Stacked.init(Block, Module)(named=initial_named, array=jax.numpy.ones(Block.size), static=1)
+
+    x = hax.random.uniform(jax.random.PRNGKey(1), (E,))
+    y, intermediates = m.scan(x, key=jax.random.split(jax.random.PRNGKey(2), Block.size))
+
+    assert y.axes == (E,)
+    assert intermediates.axes == (Block, E)

haliax-1.4.dev313/tests/test_tree_util.py

@@ -0,0 +1,91 @@
+import dataclasses
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+from chex import assert_trees_all_close
+
+import haliax as hax
+import haliax.tree_util as htu
+from haliax import Axis
+
+
+def test_resize_axis():
+
+    A = hax.Axis("A", 10)
+    B = hax.Axis("B", 20)
+    C = hax.Axis("C", 30)
+
+    class Module(eqx.Module):
+        name1: hax.NamedArray
+        name2: hax.NamedArray
+        name3: hax.NamedArray
+
+    module = Module(
+        name1=hax.random.normal(jax.random.PRNGKey(0), (B, A, C)),
+        name2=hax.zeros((B, C)),
+        name3=hax.zeros((Axis("A", 20),)),
+    )
+
+    NewA = A.resize(15)
+
+    module2 = htu.resize_axis(module, "A", 15, key=jax.random.PRNGKey(1))
+
+    assert module2.name1.axes == (B, NewA, C)
+    assert module2.name2.axes == (B, C)
+    assert module2.name3.axes == (NewA,)
+
+    # we don't mess with the mean or std of the original array too much
+    assert jnp.allclose(module2.name1.mean(), module.name1.mean(), rtol=1e-1, atol=1e-2)
+
+
+def test_scan_aware_tree_map():
+    Embed = hax.Axis("embed", 10)
+    Up = hax.Axis("up", 20)
+    Block = hax.Axis("block", 4)
+
+    class Module(eqx.Module):
+        up: hax.nn.Linear
+        down: hax.nn.Linear
+
+        def __call__(self, x, *, key):
+            return self.down(self.up(x), key=key)
+
+        @staticmethod
+        def init(layer_idx, *, key):
+            k1, k2 = jax.random.split(key)
+            up = hax.nn.Linear.init(Embed, Up, key=k1)
+            down = hax.nn.Linear.init(Up, Embed, key=k2)
+
+            up = dataclasses.replace(up, weight=up.weight + layer_idx)  # type: ignore
+            down = dataclasses.replace(down, weight=down.weight + layer_idx)  # type: ignore
+
+            return Module(up=up, down=down)
+
+    class Model(eqx.Module):
+        layers: hax.nn.Stacked[eqx.Module]
+
+        def __call__(self, x, *, key):
+            return self.layers.fold(x, key=jax.random.split(key, self.layers.Block.size))
+
+        @staticmethod
+        def init(Layers, *, key):
+            stack = hax.nn.Stacked.init(Layers, Module)(
+                layer_idx=hax.arange(Layers), key=jax.random.split(key, Layers.size)
+            )
+            return Model(layers=stack)
+
+    model = Model.init(Block, key=jax.random.PRNGKey(0))
+
+    def transform_linear(x):
+        if not isinstance(x, hax.nn.Linear):
+            return x
+
+        # do something that distinguishes doing weights jointly from independently
+        new_weight = x.weight - hax.mean(x.weight)
+        return dataclasses.replace(x, weight=new_weight)  # type: ignore
+
+    model2 = htu.scan_aware_tree_map(transform_linear, model, is_leaf=lambda x: isinstance(x, hax.nn.Linear))
+    model3 = htu.tree_map(transform_linear, model, is_leaf=lambda x: isinstance(x, hax.nn.Linear))
+
+    assert hax.all(model2.layers.stacked.up.weight != model3.layers.stacked.up.weight)

haliax-1.4.dev308/src/haliax/__about__.py

@@ -1 +0,0 @@
-__version__ = "1.4.dev308"

haliax-1.4.dev308/tests/test_tree_util.py

@@ -1,36 +0,0 @@
-import equinox as eqx
-import jax
-import jax.numpy as jnp
-
-import haliax as hax
-import haliax.tree_util as htu
-from haliax import Axis
-
-
-def test_resize_axis():
-
-    A = hax.Axis("A", 10)
-    B = hax.Axis("B", 20)
-    C = hax.Axis("C", 30)
-
-    class Module(eqx.Module):
-        name1: hax.NamedArray
-        name2: hax.NamedArray
-        name3: hax.NamedArray
-
-    module = Module(
-        name1=hax.random.normal(jax.random.PRNGKey(0), (B, A, C)),
-        name2=hax.zeros((B, C)),
-        name3=hax.zeros((Axis("A", 20),)),
-    )
-
-    NewA = A.resize(15)
-
-    module2 = htu.resize_axis(module, "A", 15, key=jax.random.PRNGKey(1))
-
-    assert module2.name1.axes == (B, NewA, C)
-    assert module2.name2.axes == (B, C)
-    assert module2.name3.axes == (NewA,)
-
-    # we don't mess with the mean or std of the original array too much
-    assert jnp.allclose(module2.name1.mean(), module.name1.mean(), rtol=1e-1, atol=1e-2)