haliax 1.4.dev308__tar.gz → 1.4.dev310__tar.gz

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

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

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

{haliax-1.4.dev308 → haliax-1.4.dev310}/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.dev310}/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.dev310/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)