haliax 1.4.dev375__tar.gz → 1.4.dev379__tar.gz

{haliax-1.4.dev375 → haliax-1.4.dev379}/.github/workflows/run_quick_levanter_tests.yaml

@@ -9,7 +9,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [3.10, 3.11]
+        python-version: ["3.10", "3.11"]
     steps:
       - name: Checkout repository
         uses: actions/checkout@v3
@@ -39,10 +39,11 @@ jobs:
         run: |
           uv pip install -e .[tests]
           uv pip install -e ../haliax
+          uv pip install pytest
 
       - name: Run pytest suite
         env:
           XLA_FLAGS: --xla_force_host_platform_device_count=8
           VIRTUAL_ENV: ${{ github.workspace }}/.venv
         run: |
-          pytest tests -m "not entry and not slow"
+          uv run pytest tests -m "not entry and not slow"

{haliax-1.4.dev375 → haliax-1.4.dev379}/PKG-INFO

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

@@ -458,7 +458,9 @@ We also provide a way to create a sequence of layers that can be applied to a se
 same interface as [haliax.nn.Stacked][], but with a different implementation. This is the [haliax.nn.BlockSeq][] module.
 BlockSeq implements those for loops directly, rather than using [haliax.fold][] or [haliax.scan][].
 
-[haliax.nn.scan.BlockFoldable][] is an interface that both [haliax.nn.Stacked][] and [haliax.nn.BlockSeq][] implement.
+[haliax.nn.scan.BlockFoldable][] is an interface that both [haliax.nn.Stacked][] and [haliax.nn.BlockSeq][] implement. It
+exposes the usual ``fold`` and ``scan`` methods as well as helpers ``fold_via`` and ``scan_via`` which return
+callables that perform the respective operations using a custom block function.
 
 ## API
 

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

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

{haliax-1.4.dev375 → haliax-1.4.dev379}/src/haliax/core.py

@@ -241,15 +241,15 @@ class NamedArray(metaclass=NamedArrayMeta):
             raise ValueError(f"Expected scalar, got {self.array.ndim}-dimensional array")
         return self.array
 
-    def __jax_array__(self):
-        if self.ndim == 0:
-            return self.array
-        else:
-            raise ValueError(
-                "Only scalar NamedArrays can be implicitly converted to jax arrays, but "
-                f"got {self.shape} array. This error typically occurs when you pass a "
-                "NamedArray to a plain jax.numpy function. Please use `x.array` instead."
-            )
+    # def __jax_array__(self):
+    #     if self.ndim == 0:
+    #         return self.array
+    #     else:
+    #         raise ValueError(
+    #             "Only scalar NamedArrays can be implicitly converted to jax arrays, but "
+    #             f"got {self.shape} array. This error typically occurs when you pass a "
+    #             "NamedArray to a plain jax.numpy function. Please use `x.array` instead."
+    #         )
 
     @ft.cached_property
     def shape(self) -> Dict[str, int]:
@@ -1787,12 +1787,22 @@ def broadcast_arrays_and_return_axes(
     """
     if len(arrays) == 0:
         return ((), ())
+    elif len(arrays) == 1:
+        a = arrays[0]
+        if a is None:
+            return (None,), ()
+        if isinstance(a, NamedArray):
+            return (a,), a.axes
+        return (named(jnp.asarray(a), ()), ), ()
 
     # sort the arrays by size, so that we use the biggest ones to broadcast the others
     # need to hold on to the order so we can return the arrays in the same order
     actual_arrays = [x for x in arrays if isinstance(x, NamedArray)]
     size_order = sorted(range(len(actual_arrays)), key=lambda i: actual_arrays[i].size, reverse=True)
     all_axes = [actual_arrays[i].axes for i in size_order]
+    if len(all_axes) == 0:
+        return (arrays, ())
+
     full_axes = ft.reduce(lambda a, b: _broadcast_axes(a, b, require_subset) if a is not None else None, all_axes)  # type: ignore
     if full_axes is None:
         raise ValueError(f"Cannot broadcast arrays {arrays}: no subset relationship")

{haliax-1.4.dev375 → haliax-1.4.dev379}/src/haliax/nn/attention.py

@@ -155,7 +155,7 @@ def combine_masks_or(mask1: Optional[NamedArray], mask2: Optional[NamedArray]) -
     return mask1 | mask2.broadcast_axis(mask1.axes)
 
 
-def causal_mask(QPos: Axis, KPos: Axis, q_start: int = 0, k_start: int = 0) -> NamedArray:
+def causal_mask(QPos: Axis, KPos: Axis, q_start: int | NamedArray = 0, k_start: int  | NamedArray= 0) -> NamedArray:
     """
     Creates a materialized causal mask for attention.
 
@@ -163,7 +163,18 @@ def causal_mask(QPos: Axis, KPos: Axis, q_start: int = 0, k_start: int = 0) -> N
     :param KPos: Axis of key sequence length
     :return: NamedArray of shape (QPos, KPos)
     """
-    return haliax.arange(QPos, start=q_start) >= haliax.arange(KPos, start=k_start).broadcast_axis(QPos)
+    # if q_start is a named array, we vmap the arange
+    if isinstance(q_start, NamedArray):
+        q_range = haliax.vmap(haliax.arange, q_start.axes)(QPos, start=q_start)
+    else:
+        q_range = haliax.arange(QPos, start=q_start)
+
+    if isinstance(k_start, NamedArray):
+        k_range = haliax.vmap(haliax.arange, k_start.axes)(KPos, start=k_start)
+    else:
+        k_range = haliax.arange(KPos, start=k_start)
+
+    return q_range >= k_range.broadcast_axis(QPos)
 
 
 def prefix_lm_mask(QSeqLen: Axis, KSeqLen: Axis, prefix_len: int, q_start: int = 0, k_start: int = 0) -> NamedArray:

{haliax-1.4.dev375 → haliax-1.4.dev379}/src/haliax/nn/dropout.py

@@ -61,8 +61,8 @@ class Dropout(eqx.Module):
     """
 
     # key difference from equinox: these are static fields
-    pdrop: float = eqx.static_field()
-    broadcast_axes: Optional[AxisSpec] = eqx.static_field()
+    pdrop: float = eqx.field(static=True)
+    broadcast_axes: Optional[AxisSpec] = eqx.field(static=True)
     inference: bool = False  # note: not static
 
     def __init__(

{haliax-1.4.dev375 → haliax-1.4.dev379}/src/haliax/nn/embedding.py

@@ -17,8 +17,8 @@ class Embedding(eqx.Module):
     weight: NamedArray
 
     # axes
-    Vocab: Axis = eqx.static_field()
-    Embed: AxisSpec = eqx.static_field()
+    Vocab: Axis = eqx.field(static=True)
+    Embed: AxisSpec = eqx.field(static=True)
 
     @staticmethod
     def init(Vocab: Axis, Embed: AxisSpec, *, init_scale: float = 1, key, initializer_range: Optional[float] = None):

{haliax-1.4.dev375 → haliax-1.4.dev379}/src/haliax/nn/linear.py

@@ -28,8 +28,8 @@ class Linear(ModuleWithStateDictSerialization):
     weight: NamedArray
     bias: Optional[NamedArray]
 
-    In: AxisSpec = eqx.static_field()
-    Out: AxisSpec = eqx.static_field()
+    In: AxisSpec = eqx.field(static=True)
+    Out: AxisSpec = eqx.field(static=True)
     dot_general: DotGeneralOp = eqx.field(default_factory=DotGeneralOp.default)
 
     @staticmethod
@@ -140,9 +140,9 @@ class MoELinear(eqx.Module):
     weight: NamedArray
     bias: Optional[NamedArray]
 
-    Experts: AxisSpec = eqx.static_field()
-    In: Axis = eqx.static_field()
-    Out: Axis = eqx.static_field()
+    Experts: AxisSpec = eqx.field(static=True)
+    In: Axis = eqx.field(static=True)
+    Out: Axis = eqx.field(static=True)
     # TODO: support quanitization for ragged_dot?
     # dot_general: DotGeneralOp = eqx.field(default_factory=DotGeneralOp.default)
 

{haliax-1.4.dev375 → haliax-1.4.dev379}/src/haliax/nn/normalization.py

@@ -20,10 +20,10 @@ A = TypeVar("A", Scalar, NamedArray, jnp.ndarray)
 
 
 class LayerNormBase(ModuleWithStateDictSerialization):
-    axis: AxisSpec = eqx.static_field()
+    axis: AxisSpec = eqx.field(static=True)
     weight: Optional[NamedArray]
     bias: Optional[NamedArray]
-    eps: float = eqx.static_field(default=1e-5)
+    eps: float = eqx.field(default=1e-5, static=True)
     dtype: Optional[jnp.dtype] = eqx.field(default=None, static=True)
 
     @abstractmethod

{haliax-1.4.dev375 → haliax-1.4.dev379}/src/haliax/nn/scan.py

@@ -2,7 +2,8 @@ import dataclasses
 import functools
 import re
 import warnings
-from typing import Any, Callable, Dict, Generic, Optional, Protocol, Sequence, Type, TypeVar, cast, overload, ParamSpec
+from typing import Any, Callable, Concatenate, Dict, Generic, Optional, Protocol, Sequence, Type, TypeVar, cast, \
+    overload, ParamSpec
 
 import equinox as eqx
 import jax
@@ -27,13 +28,18 @@ 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 FoldFunction(Protocol[M_contra, P, CarryT]):
+    def __call__(self, module: M_contra, carry: CarryT, *args: P.args, **kwargs: P.kwargs) -> CarryT:
         ...
 
 
-class ScanFunction(Protocol[M_contra, CarryT, OutputT_co]):
-    def __call__(self, module: M_contra, carry: CarryT) -> tuple[CarryT, OutputT_co]:
+class ScanFunction(Protocol[M_contra, CarryT, P, OutputT_co]):
+    def __call__(self, module: M_contra, carry: CarryT, *args: P.args, **kwargs: P.kwargs) -> tuple[CarryT, OutputT_co]:
+        ...
+
+
+class VmapFunction(Protocol[M_contra, P, OutputT_co]):
+    def __call__(self, module: M_contra, *args: P.args, **kwargs: P.kwargs) -> OutputT_co:
         ...
 
 
@@ -43,12 +49,13 @@ class ModuleInit(Protocol[M_co]):
 
 
 class BlockFoldable(Protocol[M]):
-    """
-    A superclass for [haliax.nn.Stacked][] and [haliax.nn.BlockSeq][] that exposes the fold and scan methods, as
-    well as a few other methods that are useful for these modules.
+    """Common interface for :class:`~haliax.nn.Stacked` and :class:`~haliax.nn.BlockSeq`.
 
-    This is a protocol, so you can use it as a type hint for a function that takes a Stacked or BlockSeq.
-    Equinox modules can't directly inherit from Protocols, but you can use it as a type hint.
+    The interface exposes the :py:meth:`fold` and :py:meth:`scan` methods along with the helper
+    methods :py:meth:`fold_via`, :py:meth:`scan_via`, and :py:meth:`vmap_via`.
+
+    This is a protocol, so you can use it as a type hint for a function that takes a ``Stacked`` or ``BlockSeq``.
+    Equinox modules can't directly inherit from ``Protocol`` classes, but you can use it as a type hint.
     """
 
     Block: Axis
@@ -70,6 +77,39 @@ class BlockFoldable(Protocol[M]):
     def fold(self, init: T, *args, **kwargs) -> T:
         ...
 
+    @overload
+    def fold_via(self, fn: FoldFunction[M, P, CarryT]) -> Callable[Concatenate[CarryT, P], CarryT]:
+        ...
+
+    @overload
+    def fold_via(self, fn: Callable[[M, CarryT], CarryT]) -> Callable[[CarryT], CarryT]:
+        ...
+
+    def fold_via(self, fn: Callable[..., CarryT]) -> Callable[Concatenate[CarryT, P], CarryT]:
+        ...
+
+    @overload
+    def scan_via(self, fn: ScanFunction[M, CarryT, P, OutputT_co]) -> Callable[Concatenate[CarryT, P], 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]]) -> Callable[P, tuple[CarryT, OutputT_co]]:
+        ...
+
+    @overload
+    def vmap_via(self, fn: VmapFunction[M, P, OutputT_co]) -> Callable[P, OutputT_co]:
+        ...
+
+    @overload
+    def vmap_via(self, fn: Callable[[M], OutputT_co]) -> Callable[[], OutputT_co]:
+        ...
+
+    def vmap_via(self, fn: Callable[..., OutputT_co]) -> Callable[..., OutputT_co]:
+        ...
+
     def unstacked(self) -> Sequence[M]:
         """
         Returns the unstacked version of this module. This is useful for logging or saving checkpoints.
@@ -90,8 +130,8 @@ class BlockSeq(ModuleWithStateDictSerialization, Generic[M]):
     """
 
     blocks: Sequence[M]
-    Block: Axis = eqx.static_field()
-    gradient_checkpointing: ScanCheckpointPolicy = eqx.static_field()
+    Block: Axis = eqx.field(static=True)
+    gradient_checkpointing: ScanCheckpointPolicy = eqx.field(static=True)
 
     @classmethod
     def init(
@@ -177,6 +217,87 @@ 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]:
+        ...
+
+    @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 sequence 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_fold(init: CarryT, *args, **kwargs) -> CarryT:
+            carry = init
+            for block in self.blocks:
+                carry = fn(block, carry, *args, **kwargs)
+                carry = tree_checkpoint_name(carry, self._carry_ckpt_name)
+            return carry
+
+        return do_fold
+
+    @overload
+    def scan_via(self, fn: ScanFunction[M, CarryT, P, OutputT_co]) -> Callable[Concatenate[CarryT, P], 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 sequence 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_scan(init: CarryT, *args, **kwargs) -> tuple[CarryT, OutputT_co]:
+            out = []
+            carry = init
+            for block in self.blocks:
+                carry, extra = fn(block, carry, *args, **kwargs)
+                carry = tree_checkpoint_name(carry, self._carry_ckpt_name)
+                extra = tree_checkpoint_name(extra, self._output_ckpt_name)
+                out.append(extra)
+
+            stacked_out = haliax.tree_util.tree_map(lambda *x: haliax.stack(self.Block, x), *out)
+            return carry, stacked_out
+
+        return do_scan
+
+    @overload
+    def vmap_via(self, fn: VmapFunction[M, P, OutputT_co]) -> Callable[P, OutputT_co]:
+        ...
+
+    @overload
+    def vmap_via(self, fn: Callable[[M], OutputT_co]) -> Callable[[], OutputT_co]:
+        ...
+
+    def vmap_via(self, fn: Callable[..., OutputT_co]) -> Callable[..., OutputT_co]:
+        """Return a function that applies each block independently using ``fn``.
+
+        ``fn`` should take a block and a carry and return an output. The
+        returned function mirrors :func:`haliax.vmap` over the block axis.
+        """
+
+        def do_vmap(init: CarryT, *args, **kwargs) -> OutputT_co:
+            # Apply fn to each block independently
+            outputs = []
+            for block in self.blocks:
+                output = fn(block, init, *args, **kwargs)
+                outputs.append(output)
+
+            # Stack the outputs
+            stacked_out = haliax.tree_util.tree_map(lambda *x: haliax.stack(self.Block, x), *outputs)
+            return stacked_out
+
+        return do_vmap
+
     def unstacked(self) -> Sequence[M]:
         return self.blocks
 
@@ -279,8 +400,8 @@ class Stacked(ModuleWithStateDictSerialization, Generic[M]):
     # TODO: we can probably make this module support pipeline parallelism, but that's a whole project in itself
 
     stacked: M
-    Block: Axis = eqx.static_field()
-    gradient_checkpointing: ScanCheckpointPolicy = eqx.static_field()
+    Block: Axis = eqx.field(static=True)
+    gradient_checkpointing: ScanCheckpointPolicy = eqx.field(static=True)
 
     @classmethod
     def init(
@@ -398,7 +519,7 @@ class Stacked(ModuleWithStateDictSerialization, Generic[M]):
         return do_fold(init, *args, **kwargs)
 
     @overload
-    def fold_via(self, fn: FoldFunction[M, CarryT]) -> Callable[[CarryT], CarryT]:
+    def fold_via(self, fn: FoldFunction[M, P, CarryT]) -> Callable[Concatenate[CarryT, P], CarryT]:
         ...
 
     @overload
@@ -412,16 +533,16 @@ class Stacked(ModuleWithStateDictSerialization, Generic[M]):
         returned function mirrors :func:`haliax.fold` over the block axis.
         """
 
-        def do_block(carry: CarryT, block: M) -> CarryT:
-            return fn(block, carry)
+        def do_block(carry: CarryT, block: M, *args, **kwargs) -> CarryT:
+            return fn(block, carry, *args, **kwargs)
 
-        def do_fold(init: CarryT) -> CarryT:
-            return haliax.fold(do_block, self.Block, remat=self.gradient_checkpointing)(init, self.stacked)
+        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 do_fold
 
     @overload
-    def scan_via(self, fn: ScanFunction[M, CarryT, OutputT_co]) -> Callable[[CarryT], tuple[CarryT, OutputT_co]]:
+    def scan_via(self, fn: ScanFunction[M, CarryT, P, OutputT_co]) -> Callable[Concatenate[CarryT, P], tuple[CarryT, OutputT_co]]:
         ...
 
     @overload
@@ -435,42 +556,47 @@ class Stacked(ModuleWithStateDictSerialization, Generic[M]):
         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_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) -> tuple[CarryT, OutputT_co]:
-            return haliax.scan(do_block, self.Block, remat=self.gradient_checkpointing)(init, self.stacked)
+
+        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 do_scan
 
-    def vmap(self, init, *extra_args, **extra_kwargs):
-        """Apply each block independently using :func:`haliax.vmap`.
+    @overload
+    def vmap_via(self, fn: VmapFunction[M, P, OutputT_co]) -> Callable[P, OutputT_co]:
+        ...
+
+    @overload
+    def vmap_via(self, fn: Callable[[M], OutputT_co]) -> Callable[[], OutputT_co]:
+        ...
+
+    def vmap_via(self, fn: Callable[..., OutputT_co]) -> Callable[..., OutputT_co]:
+        """Return a function that applies each block independently using ``fn``.
+
+        ``fn`` should take a block and a carry and return an output. The
+        returned function mirrors :func:`haliax.vmap` over the block axis.
+        """
+
+        def do_vmap(*args, **kwargs) -> OutputT_co:
+            # Create a function that captures the additional arguments
+            def do_block_with_args(block: M, *args, **kwargs) -> OutputT_co:
+                return fn(block, *args, **kwargs)
+
+            return haliax.vmap(do_block_with_args, self.Block)(self.stacked, *args, **kwargs)
 
-        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.
+        return do_vmap
+
+    def vmap(self, *extra_args, **extra_kwargs):
+        """Apply each block independently using :func:`haliax.vmap`.
 
         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)
+        return haliax.vmap(type(self.stacked).__call__, self.Block)(self.stacked, *extra_args, **extra_kwargs)
 
     @staticmethod
     def _do_block(carry, block, *extra_args, **extra_kwargs):

{haliax-1.4.dev375 → haliax-1.4.dev379}/src/haliax/ops.py

@@ -85,6 +85,8 @@ def where(
                 raise ValueError("x must be a NamedArray or scalar if y is a NamedArray")
             x = named(x, ())
         x, y = broadcast_arrays(x, y)
+        if isinstance(condition, NamedArray):
+            condition = condition.scalar()
         return jax.lax.cond(condition, lambda _: x, lambda _: y, None)
 
     condition, x, y = broadcast_arrays(condition, x, y)  # type: ignore

{haliax-1.4.dev375 → haliax-1.4.dev379}/src/haliax/wrap.py

@@ -5,6 +5,7 @@ import jax
 from haliax.core import NamedArray, _broadcast_order, broadcast_to
 
 from .axis import AxisSelection, AxisSelector, axis_spec_to_shape_dict, eliminate_axes
+from .jax_utils import is_scalarish
 
 
 def wrap_elemwise_unary(f, a, *args, **kwargs):
@@ -94,8 +95,33 @@ def wrap_elemwise_binary(op):
             b = broadcast_to(b, axes)
             return NamedArray(op(a.array, b.array), axes)
         elif isinstance(a, NamedArray):
+            # b isn't named.
+            if not is_scalarish(a):
+                if not is_scalarish(b):
+                    raise ValueError(
+                        f"Cannot apply {op.__name__} to a NamedArray and a non-scalar {type(b)}. "
+                        "Either both must be NamedArrays or b must be a scalar."
+                    )
+            else:
+                if is_scalarish(b):
+                    return NamedArray(op(a.array, b), a.axes)
+                a = a.scalar()
+                return op(a, b)
+
             return NamedArray(op(a.array, b), a.axes)
         elif isinstance(b, NamedArray):
+            if not is_scalarish(b):
+                if not is_scalarish(a):
+                    raise ValueError(
+                        f"Cannot apply {op.__name__} to a non-scalar {type(a)} and a NamedArray. "
+                        "Either both must be NamedArrays or a must be a scalar."
+                    )
+            else:
+                if is_scalarish(a):
+                    return NamedArray(op(a, b.array), b.axes)
+                b = b.scalar()
+                return op(a, b)
+
             return NamedArray(op(a, b.array), b.axes)
         else:
             return op(a, b)

{haliax-1.4.dev375 → haliax-1.4.dev379}/tests/core_test.py

@@ -34,7 +34,7 @@ def test_reduction_functions():
     m1 = NamedArray(rand_m, (Height, Width, Depth))
 
     # sum out everything
-    assert jnp.all(jnp.equal(hax.sum(m1), jnp.sum(m1.array)))
+    assert jnp.all(jnp.equal(hax.sum(m1).array, jnp.sum(m1.array)))
     # ensure it's a scalar
 
     assert jnp.all(jnp.equal(hax.sum(m1, axis=Height).array, jnp.sum(m1.array, axis=0)))
@@ -60,7 +60,7 @@ def test_reduction_functions():
     )
 
     # argmax
-    assert jnp.all(jnp.equal(hax.argmax(m1, axis=None), jnp.argmax(m1.array)))
+    assert jnp.all(jnp.equal(hax.argmax(m1, axis=None).array, jnp.argmax(m1.array)))
     assert jnp.all(jnp.equal(hax.argmax(m1, axis=Height).array, jnp.argmax(m1.array, axis=0)))
 
 
@@ -75,7 +75,7 @@ def test_reduction_functions_with_where():
     jmask = m1.array > 0.5
 
     # sum out everything
-    assert jnp.all(jnp.equal(hax.sum(m1, where=mask), jnp.sum(rand_m, where=jmask)))
+    assert jnp.all(jnp.equal(hax.sum(m1, where=mask).array, jnp.sum(rand_m, where=jmask)))
     # ensure it's a scalar
 
     assert jnp.all(jnp.equal(hax.sum(m1, axis=H, where=mask).array, jnp.sum(rand_m, axis=0, where=jmask)))
@@ -507,7 +507,7 @@ def test_index():
     assert named1[{"H": slice(0, 10, 2), "W": 0}].axes == (Axis("H", 5), D)
 
     # try indexing with 3 integers: returns scalar ndarray
-    assert jnp.all(jnp.equal(named1[{"H": 0, "W": 0, "D": 0}], named1.array[0, 0, 0]))
+    assert jnp.all(jnp.equal(named1[{"H": 0, "W": 0, "D": 0}].array, named1.array[0, 0, 0]))
 
 
 def test_index_with_tracer():

{haliax-1.4.dev375 → haliax-1.4.dev379}/tests/test_dot.py

@@ -31,7 +31,7 @@ def test_dot():
     # reduce to scalar
     assert jnp.all(
         jnp.equal(
-            hax.dot(m1, m2, axis=None),
+            hax.dot(m1, m2, axis=None).array,
             jnp.einsum("ijk,kji->", m1.array, m2.array),
         )
     )

{haliax-1.4.dev375 → haliax-1.4.dev379}/tests/test_hof.py

@@ -79,7 +79,7 @@ def test_scan_doesnt_scan_scalars():
     named1 = hax.random.uniform(PRNGKey(0), (Height,))
 
     def scan_fun(acc, z, x):
-        return acc + z * x, x * z
+        return (acc + z * x).scalar(), x * z
 
     total, selected = hax.scan(scan_fun, Height)(0.0, 4.0, named1)
 

{haliax-1.4.dev375 → haliax-1.4.dev379}/tests/test_ops.py

@@ -215,14 +215,14 @@ def test_mean_respects_where():
     named1 = hax.random.uniform(PRNGKey(0), (Height, Width))
     where = hax.random.uniform(PRNGKey(1), (Height, Width)) > 0.5
 
-    assert not jnp.all(jnp.isclose(hax.mean(named1), hax.mean(named1, where=where)))
-    assert jnp.all(jnp.isclose(hax.mean(named1, where=where), jnp.mean(named1.array, where=where.array)))
+    assert not hax.all(hax.isclose(hax.mean(named1), hax.mean(named1, where=where)))
+    assert jnp.all(jnp.isclose(hax.mean(named1, where=where).array, jnp.mean(named1.array, where=where.array)))
 
     # check broadcasting
     where = hax.random.uniform(PRNGKey(2), (Height,)) > 0.5
-    assert not jnp.all(jnp.isclose(hax.mean(named1), hax.mean(named1, where=where)))
+    assert not jnp.all(jnp.isclose(hax.mean(named1).array, hax.mean(named1, where=where).array))
     assert jnp.all(
-        jnp.isclose(hax.mean(named1, where=where), jnp.mean(named1.array, where=where.array.reshape((-1, 1))))
+        jnp.isclose(hax.mean(named1, where=where).array, jnp.mean(named1.array, where=where.array.reshape((-1, 1))))
     )