haliax 1.4.dev310__tar.gz → 1.4.dev313__tar.gz

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

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

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