haliax 1.4.dev367__tar.gz → 1.4.dev369__tar.gz

haliax-1.4.dev369/.playbooks/add-types.md

@@ -0,0 +1,45 @@
+# Adding NamedArray Type Annotations
+
+This playbook explains how to migrate existing code to use the new type
+annotation helpers described in `docs/typing.md`.
+
+1. **Read the documentation**: Familiarise yourself with `docs/typing.md`.
+   It describes the `Named[...]` syntax and the dtype aware helpers in
+   `haliax.typing`.
+2. **Annotate parameters**: Replace plain `NamedArray` annotations with
+   `Named[...]` that lists the required axes.  Use ellipses or sets when the
+   exact order is flexible.
+
+   ```python
+   from haliax import Named
+
+   # old
+   def foo(x: NamedArray) -> NamedArray:
+       ...
+
+   # new
+   def foo(x: Named["batch", "embed"]) -> Named["batch", "embed"]:
+       ...
+   ```
+3. **Annotate dtypes when needed**: If the dtype matters, import symbolic dtypes
+   from `haliax.typing` (e.g. `ht.f32`, `ht.i32`).  They can be indexed just
+   like `Named`.
+
+   ```python
+   import haliax.typing as ht
+
+   def bar(x: ht.f32["batch"]):
+       ...
+   ```
+4. **Runtime validation**: Use `arr.matches_axes(...)` to check that a
+   `NamedArray` conforms to the expected axes and dtype at runtime.
+
+   ```python
+   if not arr.matches_axes(Named["batch embed ..."]):
+       raise ValueError("unexpected axes")
+   ```
+5. **Update return types**: Functions returning `NamedArray` should annotate
+   their return values using the same conventions.
+
+Following these steps will gradually port legacy code to the new
+annotation style.

{haliax-1.4.dev367 → haliax-1.4.dev369}/AGENTS.md

@@ -16,8 +16,7 @@ repository. Follow these notes when implementing new features or fixing bugs.
 
 ## 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-typing.md
 
 ## Code Style
 
@@ -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.dev367 → haliax-1.4.dev369}/PKG-INFO

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

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

{haliax-1.4.dev367 → haliax-1.4.dev369}/src/haliax/nn/scan.py

@@ -233,6 +233,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 +384,35 @@ class Stacked(ModuleWithStateDictSerialization, Generic[M]):
 
         return do_fold(init, *args, **kwargs)
 
+    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.dev367 → haliax-1.4.dev369}/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

haliax-1.4.dev367/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.dev367/src/haliax/__about__.py

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