haliax 1.4.dev441__tar.gz → 1.4.dev444__tar.gz

{haliax-1.4.dev441 → haliax-1.4.dev444}/PKG-INFO

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

@@ -3,4 +3,4 @@
 # SPDX-License-Identifier: Apache-2.0
 
 
-__version__ = "1.4.dev441"
+__version__ = "1.4.dev444"

{haliax-1.4.dev441 → haliax-1.4.dev444}/src/haliax/nn/embedding.py

@@ -24,7 +24,12 @@ class Embedding(eqx.Module, ReparamEnabled):
     # axes
     Vocab: Axis = eqx.field(static=True)
     Embed: AxisSpec = eqx.field(static=True)
-    reparam: AbstractEmbeddingReparam = eqx.field(static=True)
+
+    _reparam_cls: type[AbstractEmbeddingReparam] = eqx.field(static=True, default=EmbeddingStandardParam)
+
+    @property
+    def reparam(self) -> AbstractEmbeddingReparam:
+        return self._reparam_cls(self.Embed, self.Vocab)
 
     @staticmethod
     def init(
@@ -61,7 +66,7 @@ class Embedding(eqx.Module, ReparamEnabled):
         weight = hax.random.truncated_normal(key, all_axes, -3, 3) * (
             init_scale * reparam_cls.init_scale(Vocab, Embed)
         )
-        return Embedding(weight=weight, Vocab=Vocab, Embed=Embed, reparam=reparam_cls(Embed, Vocab))
+        return Embedding(weight=weight, Vocab=Vocab, Embed=Embed, _reparam_cls=reparam_cls)
 
     def __call__(self, input_ids: NamedArray, *, key: PRNGKeyArray | None = None):
         """Alias for `embed`. key is ignored."""

{haliax-1.4.dev441 → haliax-1.4.dev444}/src/haliax/nn/linear.py

@@ -45,9 +45,14 @@ class Linear(ModuleWithStateDictSerialization, ReparamEnabled):
 
     In: AxisSpec = eqx.field(static=True)
     Out: AxisSpec = eqx.field(static=True)
-    reparam: AbstractLinearReparam = eqx.field(static=True)
     dot_general: DotGeneralOp = eqx.field(default_factory=DotGeneralOp.default)
 
+    _reparam_cls: type[AbstractLinearReparam] = eqx.field(static=True, default=LinearStandardParam)
+
+    @property
+    def reparam(self) -> AbstractLinearReparam:
+        return self._reparam_cls(self.In, self.Out)
+
     @staticmethod
     def init(
         In: AxisSpec,
@@ -77,7 +82,7 @@ class Linear(ModuleWithStateDictSerialization, ReparamEnabled):
         if dot_general is None:
             dot_general = DotGeneralOp.default()
 
-        return Linear(weight, bias, In, Out, dot_general=dot_general, reparam=reparam_cls(In, Out))
+        return Linear(weight, bias, In, Out, dot_general=dot_general, _reparam_cls=reparam_cls)
 
     @named_call
     def __call__(self, inputs, *, key: PRNGKeyArray | None = None):

{haliax-1.4.dev441 → haliax-1.4.dev444}/src/haliax/nn/scan.py

@@ -62,6 +62,22 @@ class ModuleInit(Protocol[M_co]):
     def __call__(self, *args, **kwargs) -> M_co: ...
 
 
+def _normalize_unroll(unroll: int | bool | None, block_size: int) -> int | bool:
+    """Convert user-provided ``unroll`` values into something understood by ``jax.lax.scan``."""
+
+    if unroll is None:
+        return 1
+
+    if isinstance(unroll, bool):
+        return unroll
+
+    resolved = int(unroll)
+    if resolved < 1:
+        raise ValueError(f"unroll must be >= 1; got {resolved}.")
+
+    return resolved
+
+
 class BlockFoldable(Protocol[M]):
     """Common interface for :class:`~haliax.nn.Stacked` and :class:`~haliax.nn.BlockSeq`.
 
@@ -84,29 +100,40 @@ class BlockFoldable(Protocol[M]):
         prevent_cse: bool = False,
     ) -> ModuleInit[S]: ...
 
-    def scan(self, init: T, *extra_args, **extra_kwargs): ...
+    def scan(self, init: T, *extra_args, unroll: int | bool | None = None, **extra_kwargs): ...
 
-    def fold(self, init: T, *args, **kwargs) -> T: ...
+    def fold(self, init: T, *args, unroll: int | bool | None = None, **kwargs) -> T: ...
 
     @overload
-    def fold_via(self, fn: FoldFunction[M, P, CarryT]) -> Callable[Concatenate[CarryT, P], CarryT]: ...
+    def fold_via(
+        self, fn: FoldFunction[M, P, CarryT], *, unroll: int | bool | None = None
+    ) -> Callable[Concatenate[CarryT, P], CarryT]: ...
 
     @overload
-    def fold_via(self, fn: Callable[[M, CarryT], CarryT]) -> Callable[[CarryT], CarryT]: ...
+    def fold_via(
+        self, fn: Callable[[M, CarryT], CarryT], *, unroll: int | bool | None = None
+    ) -> Callable[[CarryT], CarryT]: ...
 
-    def fold_via(self, fn: Callable[..., CarryT]) -> Callable[Concatenate[CarryT, P], CarryT]: ...
+    def fold_via(
+        self, fn: Callable[..., CarryT], *, unroll: int | bool | None = None
+    ) -> Callable[Concatenate[CarryT, P], CarryT]: ...
 
     @overload
     def scan_via(
-        self, fn: ScanFunction[M, CarryT, P, OutputT_co]
+        self, fn: ScanFunction[M, CarryT, P, OutputT_co], *, unroll: int | bool | None = None
     ) -> Callable[Concatenate[CarryT, P], tuple[CarryT, OutputT_co]]: ...
 
     @overload
     def scan_via(
-        self, fn: Callable[[M, CarryT], tuple[CarryT, OutputT_co]]
+        self,
+        fn: Callable[[M, CarryT], tuple[CarryT, OutputT_co]],
+        *,
+        unroll: int | bool | None = None,
     ) -> Callable[[CarryT], tuple[CarryT, OutputT_co]]: ...
 
-    def scan_via(self, fn: Callable[..., tuple[CarryT, OutputT_co]]) -> Callable[P, tuple[CarryT, OutputT_co]]: ...
+    def scan_via(
+        self, fn: Callable[..., tuple[CarryT, OutputT_co]], *, unroll: int | bool | None = None
+    ) -> Callable[P, tuple[CarryT, OutputT_co]]: ...
 
     @overload
     def vmap_via(self, fn: VmapFunction[M, P, OutputT_co]) -> Callable[P, OutputT_co]: ...
@@ -181,15 +208,15 @@ class BlockSeq(ModuleWithStateDictSerialization, Generic[M]):
 
         return fn
 
-    def scan(self, init: T, *extra_args, **extra_kwargs):
+    def scan(self, init: T, *extra_args, unroll: int | bool | None = None, **extra_kwargs):
         def do_scan(init, *extra_args, **extra_kwargs):
             out = []
             carry = init
 
             for i, block in enumerate(self.blocks):
-
                 (block_args, block_kwargs) = haliax.tree_util.tree_map(
-                    functools.partial(BlockSeq._slice_out, self.Block, i), (extra_args, extra_kwargs)
+                    functools.partial(BlockSeq._slice_out, self.Block, i),
+                    (extra_args, extra_kwargs),
                 )
 
                 block_result = block(carry, *block_args, **block_kwargs)
@@ -210,12 +237,13 @@ class BlockSeq(ModuleWithStateDictSerialization, Generic[M]):
 
         return do_scan(init, *extra_args, **extra_kwargs)
 
-    def fold(self, init: T, *args, **kwargs) -> T:
+    def fold(self, init: T, *args, unroll: int | bool | None = None, **kwargs) -> T:
         def do_fold(init, *args, **kwargs):
             carry = init
             for i, block in enumerate(self.blocks):
                 (block_args, block_kwargs) = haliax.tree_util.tree_map(
-                    functools.partial(BlockSeq._slice_out, self.Block, i), (args, kwargs)
+                    functools.partial(BlockSeq._slice_out, self.Block, i),
+                    (args, kwargs),
                 )
                 carry = block(carry, *block_args, **block_kwargs)
                 carry = tree_checkpoint_name(carry, self._carry_ckpt_name)
@@ -224,12 +252,16 @@ class BlockSeq(ModuleWithStateDictSerialization, Generic[M]):
         return do_fold(init, *args, **kwargs)
 
     @overload
-    def fold_via(self, fn: FoldFunction[M, P, CarryT]) -> Callable[Concatenate[CarryT, P], CarryT]: ...
+    def fold_via(
+        self, fn: FoldFunction[M, P, CarryT], *, unroll: int | bool | None = None
+    ) -> Callable[Concatenate[CarryT, P], CarryT]: ...
 
     @overload
-    def fold_via(self, fn: Callable[[M, CarryT], CarryT]) -> Callable[[CarryT], CarryT]: ...
+    def fold_via(
+        self, fn: Callable[[M, CarryT], CarryT], *, unroll: int | bool | None = None
+    ) -> Callable[[CarryT], CarryT]: ...
 
-    def fold_via(self, fn: Callable[..., CarryT]):
+    def fold_via(self, fn: Callable[..., CarryT], *, unroll: int | bool | None = None):
         """Return a function that folds over the sequence using ``fn``.
 
         ``fn`` should take a block and a carry and return a new carry. The
@@ -247,15 +279,18 @@ class BlockSeq(ModuleWithStateDictSerialization, Generic[M]):
 
     @overload
     def scan_via(
-        self, fn: ScanFunction[M, CarryT, P, OutputT_co]
+        self, fn: ScanFunction[M, CarryT, P, OutputT_co], *, unroll: int | bool | None = None
     ) -> Callable[Concatenate[CarryT, P], tuple[CarryT, OutputT_co]]: ...
 
     @overload
     def scan_via(
-        self, fn: Callable[[M, CarryT], tuple[CarryT, OutputT_co]]
+        self,
+        fn: Callable[[M, CarryT], tuple[CarryT, OutputT_co]],
+        *,
+        unroll: int | bool | None = None,
     ) -> Callable[[CarryT], tuple[CarryT, OutputT_co]]: ...
 
-    def scan_via(self, fn: Callable[..., tuple[CarryT, OutputT_co]]):
+    def scan_via(self, fn: Callable[..., tuple[CarryT, OutputT_co]], *, unroll: int | bool | None = None):
         """Return a function that scans over the sequence using ``fn``.
 
         ``fn`` should take a block and a carry and return ``(carry, output)``.
@@ -447,7 +482,7 @@ class Stacked(ModuleWithStateDictSerialization, Generic[M]):
 
         return fn
 
-    def scan(self, init, *extra_args, **extra_kwargs):
+    def scan(self, init, *extra_args, unroll: int | bool | None = None, **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 outputs as
@@ -475,19 +510,24 @@ class Stacked(ModuleWithStateDictSerialization, Generic[M]):
 
         """
 
+        resolved_unroll = _normalize_unroll(unroll, self.Block.size)
+
         def do_block(carry, block, *args, **kwargs):
             carry, out = block(carry, *args, **kwargs)
             return carry, out
 
         def do_scan(init, *extra_args, **extra_kwargs):
-            carry, out = haliax.scan(do_block, self.Block, remat=self.gradient_checkpointing)(
-                init, self.stacked, *extra_args, **extra_kwargs
-            )
+            carry, out = haliax.scan(
+                do_block,
+                self.Block,
+                remat=self.gradient_checkpointing,
+                unroll=resolved_unroll,
+            )(init, self.stacked, *extra_args, **extra_kwargs)
             return carry, out
 
         return do_scan(init, *extra_args, **extra_kwargs)
 
-    def fold(self, init, *args, **kwargs):
+    def fold(self, init, *args, unroll: int | bool | None = None, **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.
@@ -510,66 +550,88 @@ class Stacked(ModuleWithStateDictSerialization, Generic[M]):
 
         """
 
+        resolved_unroll = _normalize_unroll(unroll, self.Block.size)
+
         def do_block(carry, block, *args, **kwargs):
             carry = block(carry, *args, **kwargs)
             return carry
 
         def do_fold(init, *extra_args, **extra_kwargs):
-            carry = haliax.fold(do_block, self.Block, remat=self.gradient_checkpointing)(
-                init, self.stacked, *extra_args, **extra_kwargs
-            )
+            carry = haliax.fold(
+                do_block,
+                self.Block,
+                remat=self.gradient_checkpointing,
+                unroll=resolved_unroll,
+            )(init, self.stacked, *extra_args, **extra_kwargs)
             return carry
 
         return do_fold(init, *args, **kwargs)
 
     @overload
-    def fold_via(self, fn: FoldFunction[M, P, CarryT]) -> Callable[Concatenate[CarryT, P], CarryT]: ...
+    def fold_via(
+        self, fn: FoldFunction[M, P, CarryT], *, unroll: int | bool | None = None
+    ) -> Callable[Concatenate[CarryT, P], CarryT]: ...
 
     @overload
-    def fold_via(self, fn: Callable[[M, CarryT], CarryT]) -> Callable[[CarryT], CarryT]: ...
+    def fold_via(
+        self, fn: Callable[[M, CarryT], CarryT], *, unroll: int | bool | None = None
+    ) -> Callable[[CarryT], CarryT]: ...
 
-    def fold_via(self, fn: Callable[..., CarryT]):
+    def fold_via(self, fn: Callable[..., CarryT], *, unroll: int | bool | None = None):
         """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.
         """
 
+        resolved_unroll = _normalize_unroll(unroll, self.Block.size)
+
         def do_block(carry: CarryT, block: M, *args, **kwargs) -> CarryT:
             return fn(block, carry, *args, **kwargs)
 
         def do_fold(init: CarryT, *args, **kwargs) -> CarryT:
-            return haliax.fold(do_block, self.Block, remat=self.gradient_checkpointing)(
-                init, self.stacked, *args, **kwargs
-            )
+            return haliax.fold(
+                do_block,
+                self.Block,
+                remat=self.gradient_checkpointing,
+                unroll=resolved_unroll,
+            )(init, self.stacked, *args, **kwargs)
 
         return do_fold
 
     @overload
     def scan_via(
-        self, fn: ScanFunction[M, CarryT, P, OutputT_co]
+        self, fn: ScanFunction[M, CarryT, P, OutputT_co], *, unroll: int | bool | None = None
     ) -> Callable[Concatenate[CarryT, P], tuple[CarryT, OutputT_co]]: ...
 
     @overload
     def scan_via(
-        self, fn: Callable[[M, CarryT], tuple[CarryT, OutputT_co]]
+        self,
+        fn: Callable[[M, CarryT], tuple[CarryT, OutputT_co]],
+        *,
+        unroll: int | bool | None = None,
     ) -> Callable[[CarryT], tuple[CarryT, OutputT_co]]: ...
 
-    def scan_via(self, fn: Callable[..., tuple[CarryT, OutputT_co]]):
+    def scan_via(self, fn: Callable[..., tuple[CarryT, OutputT_co]], *, unroll: int | bool | None = None):
         """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.
         """
 
+        resolved_unroll = _normalize_unroll(unroll, self.Block.size)
+
         def do_block(carry: CarryT, block: M, *args, **kwargs) -> tuple[CarryT, OutputT_co]:
             carry, output = fn(block, carry, *args, **kwargs)
             return carry, output
 
         def do_scan(init: CarryT, *args, **kwargs) -> tuple[CarryT, OutputT_co]:
-            return haliax.scan(do_block, self.Block, remat=self.gradient_checkpointing)(
-                init, self.stacked, *args, **kwargs
-            )
+            return haliax.scan(
+                do_block,
+                self.Block,
+                remat=self.gradient_checkpointing,
+                unroll=resolved_unroll,
+            )(init, self.stacked, *args, **kwargs)
 
         return do_scan
 
@@ -707,7 +769,7 @@ def _unstack_state_dict(state_dict: StateDict, prefix: str | None = None) -> Sta
     for k, v in state_dict.items():
         if k.startswith(prefix) and is_jax_or_hax_array_like(v):
             for i, v_i in enumerate(v):
-                new_dict[f"{prefix}{i}.{k[len(prefix):]}"] = v_i
+                new_dict[f"{prefix}{i}.{k[len(prefix) :]}"] = v_i
         else:
             new_dict[k] = v
 

{haliax-1.4.dev441 → haliax-1.4.dev444}/tests/test_mup_embedding.py

@@ -31,7 +31,7 @@ def test_mup_embedding_unembedding_scale():
     Embed = (hax.Axis("E", 3),)
 
     weight = hax.ones(hax.concat_axis_specs(Vocab, Embed))
-    layer = Embedding(weight=weight, Vocab=Vocab, Embed=Embed, reparam=EmbeddingMup(Embed, Vocab))
+    layer = Embedding(weight=weight, Vocab=Vocab, Embed=Embed, _reparam_cls=EmbeddingMup)
 
     scale = layer.reparam.unembed_active_scale
     assert scale == pytest.approx(1.0 / hax.axis_size(Embed))

{haliax-1.4.dev441 → haliax-1.4.dev444}/tests/test_mup_linear.py

@@ -11,7 +11,12 @@ import pytest
 
 import haliax as hax
 from haliax.nn import Linear
-from haliax.nn.mup import InputLinearMup, LinearStandardParam, HiddenLinearMup, OutputLinearMup
+from haliax.nn.mup import (
+    InputLinearMup,
+    LinearStandardParam,
+    HiddenLinearMup,
+    OutputLinearMup,
+)
 
 
 @pytest.mark.parametrize("out_first", [True, False])
@@ -37,8 +42,8 @@ def test_mup_linear_call_matches_linear():
     weight = hax.ones(hax.concat_axis_specs(Out, In)) * 0.5
     bias = hax.full(Out, 0.25)
 
-    linear = Linear(weight, bias, In, Out, reparam=LinearStandardParam(In, Out))
-    mup = Linear(weight, bias, In, Out, reparam=InputLinearMup(In, Out))
+    linear = Linear(weight, bias, In, Out, _reparam_cls=LinearStandardParam)
+    mup = Linear(weight, bias, In, Out, _reparam_cls=InputLinearMup)
 
     inputs = hax.full(hax.concat_axis_specs(Batch, In), 2.0)
 
@@ -109,8 +114,8 @@ def test_input_linear_behaves_like_base_linear():
     weight = hax.ones((Out, In)) * 0.1
     bias = hax.zeros(Out)
 
-    linear = Linear(weight, bias, In, Out, reparam=LinearStandardParam(In, Out))
-    input_linear = Linear(weight, bias, In, Out, reparam=InputLinearMup(In, Out))
+    linear = Linear(weight, bias, In, Out, _reparam_cls=LinearStandardParam)
+    input_linear = Linear(weight, bias, In, Out, _reparam_cls=InputLinearMup)
 
     inputs = hax.random.normal(jrandom.PRNGKey(5), (Batch, In))
 

{haliax-1.4.dev441 → haliax-1.4.dev444}/tests/test_scan.py

@@ -6,6 +6,7 @@
 import equinox as eqx
 import jax
 import pytest
+import warnings
 from equinox import filter_grad
 
 import haliax as hax
@@ -164,6 +165,59 @@ def test_scan_with_aux_named_args():
     assert hax.all(hax.isclose(z_scan, z_seq_scan, atol=1e-5))
 
 
+def test_blockseq_unroll_never_warns():
+    class FoldModule(eqx.Module):
+        weight: hax.NamedArray
+
+        def __call__(self, carry: hax.NamedArray) -> hax.NamedArray:
+            return carry + self.weight
+
+        @staticmethod
+        def init(weight):
+            return FoldModule(weight=weight)
+
+    class ScanModule(eqx.Module):
+        weight: hax.NamedArray
+
+        def __call__(self, carry: hax.NamedArray) -> tuple[hax.NamedArray, hax.NamedArray]:
+            updated = carry + self.weight
+            return updated, updated
+
+        @staticmethod
+        def init(weight):
+            return ScanModule(weight=weight)
+
+    Block = hax.Axis("block", 3)
+    Value = hax.Axis("value", 2)
+    weights = hax.random.uniform(jax.random.PRNGKey(0), (Block, Value))
+
+    fold_seq = BlockSeq.init(Block, FoldModule)(weight=weights)
+    scan_seq = BlockSeq.init(Block, ScanModule)(weight=weights)
+
+    init_carry = hax.zeros(Value)
+
+    def fold_step(block: FoldModule, carry: hax.NamedArray) -> hax.NamedArray:
+        return block(carry)
+
+    def scan_step(block: ScanModule, carry: hax.NamedArray) -> tuple[hax.NamedArray, hax.NamedArray]:
+        return block(carry)
+
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        _ = fold_seq.fold(init_carry, unroll=2)
+        _ = fold_seq.fold(init_carry, unroll=True)
+        _ = fold_seq.fold(init_carry, unroll=False)
+        _ = fold_seq.fold_via(fold_step, unroll=True)(init_carry)
+        _ = fold_seq.fold_via(fold_step, unroll=False)(init_carry)
+        _ = scan_seq.scan(init_carry, unroll=2)
+        _ = scan_seq.scan(init_carry, unroll=True)
+        _ = scan_seq.scan(init_carry, unroll=False)
+        _ = scan_seq.scan_via(scan_step, unroll=True)(init_carry)
+        _ = scan_seq.scan_via(scan_step, unroll=False)(init_carry)
+
+    assert not caught
+
+
 def test_stacked_to_state_dict():
     class Module(eqx.Module):
         named: hax.NamedArray
@@ -239,8 +293,18 @@ E = hax.Axis("E", 10)
             [(E.size,), (Block.size, E.size), (Block.size, E.size)],
             None,
         ),
-        ("simple", ScanCheckpointPolicy(simple=True), [(E.size,), (Block.size, E.size)], None),
-        ("nested", ScanCheckpointPolicy(simple=True, nested=2), [(E.size,), (2, E.size)], None),
+        (
+            "simple",
+            ScanCheckpointPolicy(simple=True),
+            [(E.size,), (Block.size, E.size)],
+            None,
+        ),
+        (
+            "nested",
+            ScanCheckpointPolicy(simple=True, nested=2),
+            [(E.size,), (2, E.size)],
+            None,
+        ),
         (
             "sin_offload",
             ScanCheckpointPolicy(save_carries=True, offload_block_internals=["sin"]),
@@ -371,6 +435,97 @@ def test_scan_via():
     assert hax.all(hax.isclose(outs, expected_outs))
 
 
+def test_scan_via_with_unroll():
+    class Module(eqx.Module):
+        w: hax.NamedArray
+
+        def with_output(self, x):
+            out = x + self.w
+            return out, 2 * self.w
+
+        def __call__(self, carry):
+            return carry + self.w, 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,))
+
+    default_carry, default_outs = m.scan_via(Module.with_output)(x)
+    carry, outs = m.scan_via(Module.with_output, unroll=2)(x)
+
+    assert hax.all(hax.isclose(carry, default_carry))
+    assert hax.all(hax.isclose(outs, default_outs))
+
+    default_carry_direct, default_outs_direct = m.scan(x)
+    carry_direct, outs_direct = m.scan(x, unroll=2)
+
+    assert hax.all(hax.isclose(carry_direct, default_carry_direct))
+    assert hax.all(hax.isclose(outs_direct, default_outs_direct))
+
+
+def test_scan_via_with_bool_unroll(monkeypatch):
+    class Module(eqx.Module):
+        w: hax.NamedArray
+
+        def with_output(self, x):
+            out = x + self.w
+            return out, 2 * self.w
+
+        def __call__(self, carry):
+            return carry + self.w, 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,))
+
+    default_carry_via, default_outs_via = m.scan_via(Module.with_output)(x)
+    default_carry_direct, default_outs_direct = m.scan(x)
+
+    import haliax.nn.scan as hnn_scan
+
+    scan_calls: list[int | None] = []
+    original_scan = hnn_scan.haliax.scan
+
+    def wrapped_scan(*args, **kwargs):
+        scan_calls.append(kwargs.get("unroll"))
+        return original_scan(*args, **kwargs)
+
+    monkeypatch.setattr(hnn_scan.haliax, "scan", wrapped_scan)
+
+    carry_true, outs_true = m.scan_via(Module.with_output, unroll=True)(x)
+    carry_false, outs_false = m.scan_via(Module.with_output, unroll=False)(x)
+    carry_true_direct, outs_true_direct = m.scan(x, unroll=True)
+    carry_false_direct, outs_false_direct = m.scan(x, unroll=False)
+
+    assert scan_calls == [True, False, True, False]
+
+    assert hax.all(hax.isclose(carry_true, default_carry_via))
+    assert hax.all(hax.isclose(outs_true, default_outs_via))
+    assert hax.all(hax.isclose(carry_false, default_carry_via))
+    assert hax.all(hax.isclose(outs_false, default_outs_via))
+
+    assert hax.all(hax.isclose(carry_true_direct, default_carry_direct))
+    assert hax.all(hax.isclose(outs_true_direct, default_outs_direct))
+    assert hax.all(hax.isclose(carry_false_direct, default_carry_direct))
+    assert hax.all(hax.isclose(outs_false_direct, default_outs_direct))
+
+
 def test_scan_via_multi_args():
     class Module(eqx.Module):
         w: hax.NamedArray
@@ -513,6 +668,88 @@ def test_fold_via_multi_args():
     assert hax.all(hax.isclose(result, expected))
 
 
+def test_fold_via_with_unroll():
+    class Module(eqx.Module):
+        w: hax.NamedArray
+
+        def intermediate(self, x):
+            return x + 2 * self.w
+
+        def __call__(self, carry):
+            return carry + 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,))
+
+    default_result = m.fold_via(Module.intermediate)(x)
+    result_unroll = m.fold_via(Module.intermediate, unroll=2)(x)
+
+    assert hax.all(hax.isclose(result_unroll, default_result))
+
+    default_fold = m.fold(x)
+    fold_unroll = m.fold(x, unroll=2)
+
+    assert hax.all(hax.isclose(fold_unroll, default_fold))
+
+
+def test_fold_with_bool_unroll(monkeypatch):
+    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,))
+
+    default_fold = m.fold(x)
+    default_fold_via = m.fold_via(Module.intermediate)(x)
+
+    import haliax.nn.scan as hnn_scan
+
+    fold_calls: list[int | None] = []
+    original_fold = hnn_scan.haliax.fold
+
+    def wrapped_fold(*args, **kwargs):
+        fold_calls.append(kwargs.get("unroll"))
+        return original_fold(*args, **kwargs)
+
+    monkeypatch.setattr(hnn_scan.haliax, "fold", wrapped_fold)
+
+    fold_true = m.fold(x, unroll=True)
+    fold_false = m.fold(x, unroll=False)
+    via_true = m.fold_via(Module.intermediate, unroll=True)(x)
+    via_false = m.fold_via(Module.intermediate, unroll=False)(x)
+
+    assert fold_calls == [True, False, True, False]
+
+    assert hax.all(hax.isclose(fold_true, default_fold))
+    assert hax.all(hax.isclose(fold_false, default_fold))
+    assert hax.all(hax.isclose(via_true, default_fold_via))
+    assert hax.all(hax.isclose(via_false, default_fold_via))
+
+
 def test_fold_via_static_args():
     class Module(eqx.Module):
         w: hax.NamedArray