haliax 1.4.dev368__tar.gz → 1.4.dev371__tar.gz

{haliax-1.4.dev368 → haliax-1.4.dev371}/.github/workflows/run_tests.yaml

@@ -15,9 +15,8 @@ jobs:
           python-version: 3.10.11
       - name: Install dependencies
         run: |
-          python -m pip install --upgrade pip
-          pip install flake8 pytest
-          pip install -e .[dev]
+          python -m pip install uv
+          uv sync
       - name: Test with pytest
         run: |
-          XLA_FLAGS=--xla_force_host_platform_device_count=8 PYTHONPATH=tests:src:. pytest tests
+          XLA_FLAGS=--xla_force_host_platform_device_count=8 PYTHONPATH=tests:src:. uv run pytest tests

{haliax-1.4.dev368 → haliax-1.4.dev371}/AGENTS.md

@@ -12,12 +12,11 @@ repository. Follow these notes when implementing new features or fixing bugs.
   floating point tests, add that to the list.
 * **Playbooks.** Sometimes, there are repeatable tasks (e.g. porting models) for which we follow a standard set of steps.
   Please reference `.playbooks/` to see what playbooks are available, or see the list below. If you want to add a playbook
-  write a markdown doc named e.g. `.playbooks/port-models.md` and add a pointer to it in the list below.
+  write a markdown doc named e.g. `.playbooks/add-types.md` and add a pointer to it in the list below.
 
 ## Playbook
 
-- At the moment, there are no playbooks available. If you have a repeatable task that you think
-  should be documented, please create a new markdown file in `.playbooks/` and add it to the list above.
+- Adding Haliax-style tensor typing annotations are described in @.playbooks/add-types.md
 
 ## Code Style
 
@@ -45,7 +44,7 @@ repository. Follow these notes when implementing new features or fixing bugs.
 
 ## Testing
 
-* Tests are executed with `pytest`. The default workflow runs `uv run pytest tests`.
+* Tests are executed with `pytest`. The default workflow runs ` XLA_FLAGS=--xla_force_host_platform_device_count=8 PYTHONPATH=tests:src:. uv run pytest tests`.
 * In general, never relax tolerances in floating point tests unless specifically discussed with the
   team. Use `assert_allclose` with appropriate tolerances for numerical comparisons. We typically use
   1e-4 for more complex modules, and 1e-5 for simpler ones.
@@ -58,6 +57,11 @@ repository. Follow these notes when implementing new features or fixing bugs.
 
 * **Generic code**: many utilities are written with Python generics and dataclasses. Where possible,
   write reusable functions or classes that operate over TypeVars instead of hard coding concrete types.
+* **Configurations**: configuration files are dataclasses loaded via `draccus`. Keep configs
+  declarative and typed.
+* **Reproducibility**: Levanter aims for deterministic training where possible. Avoid sources of
+  nondeterminism unless explicitly required.
+* Prefer Stacked with fold or scan over writing custom loops, for better compile times and gradient checkpointing support
 
 ## Library conventions
 - Haliax revolves around `NamedArray` and explicit `Axis` objects. Prefer APIs that accept

{haliax-1.4.dev368 → haliax-1.4.dev371}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev368
+Version: 1.4.dev371
 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.dev368 → haliax-1.4.dev371}/docs/scan.md

@@ -405,6 +405,17 @@ blocks = Stacked.init(Layers, Gpt2Block)(
 Any NamedArray passed to the Stacked init will have its Layers axis (if present) vmapped over. Any
 JAX array will have its first axis vmapped over.
 
+#### Apply Blocks in Parallel with `vmap`
+
+Sometimes you may want to apply each block independently, without feeding the
+output of one block into the next.  `Stacked.vmap` does exactly that: it uses
+[`haliax.vmap`][] to broadcast the initial value to every block and evaluates
+them in parallel, returning the stack of outputs.
+
+```python
+y = stacked.vmap(x)
+```
+
 
 #### Fold Blocks vs Scan Blocks
 

haliax-1.4.dev371/docs/vmap.md

@@ -0,0 +1,9 @@
+## Vectorization
+
+
+This primitive is also used by [`Stacked.vmap`](scan.md#apply-blocks-in-parallel-with-vmap)
+to apply an entire stack of blocks in parallel.
+
+(This is a work in progress. Please contact dlwh for more information.)
+
+::: haliax.vmap

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

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

{haliax-1.4.dev368 → haliax-1.4.dev371}/src/haliax/nn/scan.py

@@ -2,7 +2,7 @@ import dataclasses
 import functools
 import re
 import warnings
-from typing import Any, Dict, Generic, Optional, Protocol, Sequence, Type, TypeVar, cast
+from typing import Any, Callable, Dict, Generic, Optional, Protocol, Sequence, Type, TypeVar, cast, overload, ParamSpec
 
 import equinox as eqx
 import jax
@@ -20,8 +20,21 @@ from ..axis import Axis
 
 M = TypeVar("M", bound=eqx.Module)
 M_co = TypeVar("M_co", bound=eqx.Module, covariant=True)
+M_contra = TypeVar("M_contra", bound=eqx.Module, contravariant=True)
 S = TypeVar("S", bound=eqx.Module)
 T = TypeVar("T")
+CarryT = TypeVar("CarryT")
+OutputT_co = TypeVar("OutputT_co", covariant=True)
+P = ParamSpec("P")
+
+class FoldFunction(Protocol[M_contra, CarryT]):
+    def __call__(self, module: M_contra, carry: CarryT) -> CarryT:
+        ...
+
+
+class ScanFunction(Protocol[M_contra, CarryT, OutputT_co]):
+    def __call__(self, module: M_contra, carry: CarryT) -> tuple[CarryT, OutputT_co]:
+        ...
 
 
 class ModuleInit(Protocol[M_co]):
@@ -233,6 +246,8 @@ class Stacked(ModuleWithStateDictSerialization, Generic[M]):
     output, while "scan" is the same as a for loop that accumulates a list of outputs as well as the final output.
 
     Stacked also supports gradient checkpointing, which is useful for very large models that don't fit in memory.
+    If your blocks are independent of each other you can instead use :py:meth:`Stacked.vmap`
+    to apply every block in parallel.
 
     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 outputs and
@@ -382,6 +397,81 @@ class Stacked(ModuleWithStateDictSerialization, Generic[M]):
 
         return do_fold(init, *args, **kwargs)
 
+    @overload
+    def fold_via(self, fn: FoldFunction[M, CarryT]) -> Callable[[CarryT], CarryT]:
+        ...
+
+    @overload
+    def fold_via(self, fn: Callable[[M, CarryT], CarryT]) -> Callable[[CarryT], CarryT]:
+        ...
+
+    def fold_via(self, fn: Callable[..., CarryT]):
+        """Return a function that folds over the stack using ``fn``.
+
+        ``fn`` should take a block and a carry and return a new carry.  The
+        returned function mirrors :func:`haliax.fold` over the block axis.
+        """
+
+        def do_block(carry: CarryT, block: M) -> CarryT:
+            return fn(block, carry)
+
+        def do_fold(init: CarryT) -> CarryT:
+            return haliax.fold(do_block, self.Block, remat=self.gradient_checkpointing)(init, self.stacked)
+
+        return do_fold
+
+    @overload
+    def scan_via(self, fn: ScanFunction[M, CarryT, OutputT_co]) -> Callable[[CarryT], tuple[CarryT, OutputT_co]]:
+        ...
+
+    @overload
+    def scan_via(self, fn: Callable[[M, CarryT], tuple[CarryT, OutputT_co]]) -> Callable[[CarryT], tuple[CarryT, OutputT_co]]:
+        ...
+
+    def scan_via(self, fn: Callable[..., tuple[CarryT, OutputT_co]]):
+        """Return a function that scans over the stack using ``fn``.
+
+        ``fn`` should take a block and a carry and return ``(carry, output)``.
+        Semantics match :func:`haliax.scan` over the block axis.
+        """
+
+        def do_block(carry: CarryT, block: M) -> tuple[CarryT, OutputT_co]:
+            return fn(block, carry)
+
+        def do_scan(init: CarryT) -> tuple[CarryT, OutputT_co]:
+            return haliax.scan(do_block, self.Block, remat=self.gradient_checkpointing)(init, self.stacked)
+
+        return do_scan
+
+    def vmap(self, init, *extra_args, **extra_kwargs):
+        """Apply each block independently using :func:`haliax.vmap`.
+
+        This maps ``init`` through every block in parallel, so each block
+        receives the same ``init`` but its own parameters.  Extra ``args`` and
+        ``kwargs`` are also mapped over the block axis by default.
+
+        Returns the stacked outputs of each block.
+        """
+
+        if haliax.is_named_array(init):
+            init = init.broadcast_axis(self.Block)
+        elif haliax.jax_utils.is_jax_array_like(init):
+            init = jnp.broadcast_to(init, (self.Block.size,) + init.shape)
+        else:
+            init = tuple(init for _ in range(self.Block.size))
+
+        arg_spec = (0, 0) + (0,) * len(extra_args)
+        kwarg_spec = {k: 0 for k in extra_kwargs}
+
+        do_vmap = haliax.vmap(
+            Stacked._do_block,
+            self.Block,
+            default=0,
+            args=arg_spec,
+            kwargs=kwarg_spec,
+        )
+        return do_vmap(init, self.stacked, *extra_args, **extra_kwargs)
+
     @staticmethod
     def _do_block(carry, block, *extra_args, **extra_kwargs):
         return block(carry, *extra_args, **extra_kwargs)

{haliax-1.4.dev368 → haliax-1.4.dev371}/tests/test_scan.py

@@ -45,6 +45,30 @@ def test_unstacked():
         assert hax.all(module.array == m.stacked.array[i])
 
 
+def test_vmap():
+    class Module(eqx.Module):
+        weight: hax.NamedArray
+
+        def __call__(self, x):
+            return x + self.weight
+
+        @staticmethod
+        def init(weight):
+            return Module(weight=weight)
+
+    Block = hax.Axis("block", 4)
+    E = hax.Axis("E", 10)
+
+    weights = hax.random.uniform(jax.random.PRNGKey(0), (Block, E))
+    m = Stacked.init(Block, Module)(weight=weights)
+
+    x = hax.random.uniform(jax.random.PRNGKey(1), (E,))
+    y = m.vmap(x)
+
+    assert y.axes == (Block, E)
+    assert hax.all(y == weights + x)
+
+
 def test_seq_and_stacked_give_same_results():
     class Module(eqx.Module):
         named: hax.NamedArray
@@ -285,3 +309,58 @@ def test_checkpoint_carries(name, policy, expected_scan_shapes, check_offloading
 
             assert target is not None, f"Could not find named value for {name}"
             assert found_saved, f"Could not find offloaded value for {name}"
+
+
+def test_fold_via():
+    class Module(eqx.Module):
+        w: hax.NamedArray
+
+        def __call__(self, x):
+            return x + self.w
+
+        def intermediate(self, x):
+            return x + 2 * self.w
+
+        @staticmethod
+        def init(named):
+            return Module(w=named)
+
+    Block = hax.Axis("block", 3)
+    E = hax.Axis("E", 5)
+
+    named = hax.random.uniform(jax.random.PRNGKey(0), (Block, E))
+    m = Stacked.init(Block, Module)(named=named)
+
+    x = hax.random.uniform(jax.random.PRNGKey(1), (E,))
+    result = m.fold_via(Module.intermediate)(x)
+
+    expected = x + 2 * hax.sum(named, Block)
+    assert hax.all(hax.isclose(result, expected))
+
+
+def test_scan_via():
+    class Module(eqx.Module):
+        w: hax.NamedArray
+
+        def with_output(self, x):
+            out = x + self.w
+            return out, 2 * self.w
+
+        @staticmethod
+        def init(named):
+            return Module(w=named)
+
+    Block = hax.Axis("block", 4)
+    E = hax.Axis("E", 6)
+
+    named = hax.random.uniform(jax.random.PRNGKey(0), (Block, E))
+    m = Stacked.init(Block, Module)(named=named)
+
+    x = hax.random.uniform(jax.random.PRNGKey(1), (E,))
+    carry, outs = m.scan_via(Module.with_output)(x)
+
+    expected_carry = x + hax.sum(named, Block)
+    expected_outs = 2 * named
+
+    assert hax.all(hax.isclose(carry, expected_carry))
+    assert hax.all(hax.isclose(outs, expected_outs))

haliax-1.4.dev368/docs/vmap.md

@@ -1,5 +0,0 @@
-## Vectorization
-
-(This is a work in progress. Please contact dlwh for more information.)
-
-::: haliax.vmap

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

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