haliax 1.4.dev369__tar.gz → 1.4.dev372__tar.gz

{haliax-1.4.dev369 → haliax-1.4.dev372}/.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.dev369 → haliax-1.4.dev372}/AGENTS.md

@@ -12,11 +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
 
-- Adding Haliax-style tensor typing annotations are described in @.playbooks/add-typing.md
+- Adding Haliax-style tensor typing annotations are described in @.playbooks/add-types.md
 
 ## Code Style
 
@@ -44,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.

{haliax-1.4.dev369 → haliax-1.4.dev372}/PKG-INFO

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

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

{haliax-1.4.dev369 → haliax-1.4.dev372}/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]):
@@ -384,6 +397,52 @@ 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`.
 

{haliax-1.4.dev369 → haliax-1.4.dev372}/tests/test_scan.py

@@ -309,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))