haliax 1.4.dev364__tar.gz → 1.4.dev366__tar.gz

{haliax-1.4.dev364 → haliax-1.4.dev366}/.pre-commit-config.yaml

@@ -2,7 +2,7 @@
 # See https://pre-commit.com/hooks.html for more hooks
 exclude: ".git"
 default_stages:
-    - commit
+    - pre-commit
 fail_fast: true
 
 repos:
@@ -16,24 +16,24 @@ repos:
     -   id: check-merge-conflict
     -   id: check-added-large-files
 
--   repo: https://github.com/psf/black
-    rev: 22.3.0
-    hooks:
-    -   id: black
-
--   repo: https://github.com/timothycrosley/isort
-    rev: 5.11.5
-    hooks:
-    -   id: isort
+- repo: https://github.com/astral-sh/ruff-pre-commit
+  rev: v0.11.10
+  hooks:
+    - id: ruff
+      args: [ --fix, --exit-non-zero-on-fix ]
 
--   repo: https://github.com/PyCQA/flake8
-    rev: 3.9.2
-    hooks:
-    -   id: flake8
-        additional_dependencies: [flake8-isort]
+#- repo: local
+#  hooks:
+#    - id: ty-check
+#      name: ty-check
+#      language: python
+#      entry: ty check
+#      pass_filenames: false
+#      args: [--python=.venv/]
+#      additional_dependencies: [ty]
 
 -   repo: https://github.com/pre-commit/mirrors-mypy
-    rev: 'v1.5.1'
+    rev: 'v1.16.1'
     hooks:
     -   id: mypy
         args: [--ignore-missing-imports, --check-untyped-defs]

haliax-1.4.dev366/AGENTS.md

@@ -0,0 +1,79 @@
+# Haliax LLM Agent Guidelines
+
+This document summarizes important conventions for contributing code or documentation to the Haliax
+repository. Follow these notes when implementing new features or fixing bugs.
+
+## General Guidelines
+
+* **Get better.** Whenever you discover something missing from these guidelines, or the requester
+  suggests a better way to do something, please update this document. The goal is to make it easier for
+  everyone to contribute and maintain the codebase. Generally speaking, you should add bullets or new sections.
+  Be sure to do this when directed to. For example, if directed that you should never relax tolerances in
+  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.
+
+## 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.
+
+## Code Style
+
+* **Python version**: the project targets Python >=3.10.
+* **Formatting and Linting**: We use `ruff` via `pre-commit`.
+* **Typing**: the code base uses `mypy` for static type checking. `mypy` is run by pre‑commit and the
+  configuration is found in `pyproject.toml`.
+* **Run `pre-commit run --all-files`** before committing. The CI workflows run the same checks.
+* **Doc Strings**: All public functions, classes, and modules should have docstrings, unless
+  their purpose is painfully obvious. Use
+  [Google style](https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings) for
+  consistency.
+* **Commenting**: Use comments to explain why something is done a certain way, especially if it is not
+  immediately obvious. Avoid commenting on every line of code; focus on the intent and purpose of
+  complex logic. Demarcating logical groups of code with comments is encouraged, unless it is better
+  to refactor the code into smaller functions or classes.
+* **Mkdocs**: We use [Mkdocs](https://www.mkdocs.org/) for documentation. The main documentation is in
+  the `docs` directory. Use Markdown for writing docs, and follow the existing structure. When linking to
+  symbols, prefer using mkdocs-style links (e.g. With a custom title: `[full.path.object2][]` or
+  `[Object 1][full.path.object1]`)
+* **Documentation**: When adding new features, ensure that the documentation is updated accordingly.
+  This includes updating the Mkdocs files and any relevant docstrings. If you add a new module or
+  significant functionality, consider adding a dedicated section in the documentation.
+
+## Testing
+
+* Tests are executed with `pytest`. The default workflow runs
+  `pytest tests -m "not entry and not slow and not ray"`.
+* 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.
+* Tests should be reasonably fast. Mark long-running tests with @pytest.mark.slow so they are excluded from the default suite.
+* Always mark tests that depend on pytorch with `@skip_if_no_torch` to ensure they are skipped
+  when PyTorch is not available. This is particularly important for tests that require PyTorch-specific
+  functionality.
+
+
+## Design Preferences
+
+* **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
+  axes or axis names rather than hard‑coding positional dimensions.
+- Utilities should be written so they work with arbitrary axis names. Avoid relying on
+  fixed axis orders when possible.
+- Use the provided modules in `haliax.nn` or Equinox when building neural network layers.
+- Type annotations can use named shapes shorthand provided in `haliax.haxtyping`: `ht.f32[NamedArray, "batch"]`
+  for a float32 array with a "batch" axis, or `ht.Float[NamedArray, "batch"]` for any floating point dtype.
+
+## Documentation
+- Public functions and modules require docstrings. If behavior is non‑obvious,
+  add examples in `docs/`.

{haliax-1.4.dev364 → haliax-1.4.dev366}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev364
+Version: 1.4.dev366
 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/
@@ -29,6 +29,7 @@ Requires-Dist: mkdocs>=1.4.3; extra == 'dev'
 Requires-Dist: mkdocstrings-python>=1.1.2; extra == 'dev'
 Requires-Dist: mkdocstrings>=0.22.0; extra == 'dev'
 Requires-Dist: mypy>=0.910; extra == 'dev'
+Requires-Dist: pre-commit; extra == 'dev'
 Requires-Dist: pygments; extra == 'dev'
 Requires-Dist: pymdown-extensions; extra == 'dev'
 Requires-Dist: pytest>=7.4.0; extra == 'dev'

{haliax-1.4.dev364 → haliax-1.4.dev366}/docs/nn.md

@@ -55,7 +55,6 @@ We don't provide an explicit attention module, but we do provide an attention fu
 
 :::haliax.nn.attention.dot_product_attention
 :::haliax.nn.attention.dot_product_attention_weights
-:::haliax.nn.attention.self_attention
 
 ### Masks
 ::: haliax.nn.attention.causal_mask

haliax-1.4.dev366/docs/typing.md

@@ -0,0 +1,100 @@
+from haliax import NamedArrayfrom haliax import NamedArray
+
+# NamedArray Type Annotations
+
+Haliax supports an extension to [`jaxtyping`](https://docs.kidger.site/jaxtyping/)
+that allows you to annotate functions and methods that take or return
+[`NamedArray`][haliax.core.NamedArray] objects. If you are familiar with
+[`jaxtyping`](https://docs.kidger.site/jaxtyping/), the syntax is very similar.
+In fact, for non-NamedArrays, it is exactly the same.
+
+```python
+from haliax import NamedArray
+import haliax.haxtyping as ht
+
+def foo(x: ht.Float[NamedArray, "batch embed ..."]):
+    ...
+```
+
+At runtime you can verify that a `NamedArray` conforms to a particular
+annotation using `matches_axes`:
+
+```python
+if not arr.matches_axes(Float[NamedArray, "batch embed ..."]):
+    raise ValueError("unexpected axes")
+```
+
+## DType-aware annotations
+
+Sometimes it is useful to express both the axes **and** the dtype in the type
+annotation.  The :mod:`haliax.typing` module defines symbolic types for all of
+JAX's common dtypes that can be indexed just like ``Named``.  In documentation
+examples we'll use ``import haliax.typing as ht``:
+
+```python
+import haliax.haxtyping as ht
+
+def foo(x: ht.f32[NamedArray, "batch"]):
+    ...
+
+def bar(x: ht.i32[NamedArray, "batch"]):
+    ...
+```
+
+For convenience the module also provides aggregate categories ``Float``,
+``Complex``, ``Int`` and ``UInt`` that match any floating point, complex,
+signed integer or unsigned integer dtype respectively:
+
+```python
+def baz(x: ht.Float[NamedArray, "batch"]):
+    ...
+```
+
+At runtime ``matches_axes`` also checks the dtype when one is present:
+
+```python
+from haliax import Axis, zeros
+import haliax.haxtyping as ht
+
+arr = zeros({"batch": 4})
+assert arr.matches_axes(ht.f32["batch"])  # dtype and axes both match
+```
+
+## FAQ
+
+### Why not use `NamedArray` directly in type annotations?
+
+Using `NamedArray` directly in type annotations doesn't work well with
+type checkers like `mypy` or `pyright`. These tools expect types to be
+subscripted with other types or forward references (which are strings).
+Using `NamedArray` directly would lead to type errors.
+
+### Why not use `jaxtyping` directly?
+
+While `jaxtyping` is a powerful library for type annotations in JAX, it does not
+support `NamedArray` objects directly. The `haliax.haxtyping` module extends
+`jaxtyping` to include `NamedArray` support, allowing you to annotate functions
+and methods that take or return `NamedArray` objects with specific axes and dtypes.
+
+### Why do I have to specify the `NamedArray` type in the annotation?
+
+I hate this, but it's the only way to get type checkers like `mypy` and `pyright` to understand that the type is
+a `NamedArray`. Underneath the hood, during type checking, `jaxtyping.Float` (and `haxtyping.Float`) are
+essentially type aliases of [`Annotated`](https://docs.python.org/3/library/typing.html#typing.Annotated)
+with the `NamedArray` type. There's no other way I could find to get type checkers to understand that the type is a
+`NamedArray` or to accept strings like `"batch embed ..."` as valid type annotations.
+
+### How do I use single axes in type annotations with flake or ruff.
+
+Like `jaxtyping`, you need to prepend a space before the axis name to use single axes in type annotations with
+flake or ruff. For example, to use a single axis named `batch`, you would write:
+
+```python
+def foo(x: ht.Float[NamedArray, " batch"]):
+    ...
+```
+
+Then suppress F722 in your linter to suppress that error.
+
+See the [jaxtyping documentation](https://docs.kidger.site/jaxtyping/faq/#flake8-or-ruff-are-throwing-an-error) for more
+details on the workaround.

{haliax-1.4.dev364 → haliax-1.4.dev366}/pyproject.toml

@@ -38,7 +38,8 @@ dev=["pytest >= 7.4.0", "mypy >= 0.910", "mkdocs >= 1.4.3", "mkdocs-material >=
     "pymdown-extensions",
     "pygments",
     "pymdown-extensions",
-    "chex>=0.1.86"
+    "chex>=0.1.86",
+    "pre-commit",
 ]
 
 
@@ -72,6 +73,11 @@ src_paths = ["src", "tests"]
 "Bug Tracker" = "https://github.com/stanford-crfm/haliax/issues/"
 "Documentation" = "https://haliax.readthedocs.io/en/latest/"
 
+[tool.ruff.lint]
+ignore = [ "E203", "E501", "W605", "F821", "E266", "F722", "E731", "E741" ]
+
+[tool.setuptools.package-data]
+"haliax" = ["*.pyi"]
 
 [tool.coverage.report]
 exclude_also = [

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

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

{haliax-1.4.dev364 → haliax-1.4.dev366}/src/haliax/__init__.py

@@ -1,4 +1,4 @@
-import typing
+import typing as t
 from typing import Optional, Sequence
 
 import jax
@@ -42,11 +42,8 @@ from .axis import (
     to_jax_shape,
 )
 from .core import (
-    Named,
     NamedArray,
-    NamedArrayAxes,
-    NamedArrayAxesSpec,
-    NamedOrNumeric,
+    NamedArrayAxes, NamedArrayAxesSpec, NamedOrNumeric,
     are_shape_checks_enabled,
     broadcast_arrays,
     broadcast_axis,
@@ -66,6 +63,7 @@ from .core import (
     unflatten_axis,
     updated_slice,
 )
+from .haxtyping import Named
 from .hof import fold, map, scan, vmap
 from .jax_utils import tree_checkpoint_name
 from .ops import clip, isclose, pad_left, trace, tril, triu, where
@@ -83,8 +81,8 @@ from .wrap import (
 )
 
 
-T = typing.TypeVar("T")
-A = typing.TypeVar("A", Scalar, NamedArray, jnp.ndarray)
+T = t.TypeVar("T")
+A = t.TypeVar("A", Scalar, NamedArray, jnp.ndarray)
 
 
 # creation routines
@@ -148,11 +146,6 @@ def arange(axis: AxisSpec, *, start=0, step=1, dtype: Optional[DTypeLike] = None
     ```
 
     """
-    from haliax.jax_utils import to_jax_shape
-    from haliax.util import ensure_tuple
-
-    # if start is a tracer, we need to be a bit cleverer since arange doesn't support tracers
-    # return NamedArray(jnp.arange(start, stop, step, dtype=dtype), (axis,))
     size = axis_size(axis)
 
     arr = jax.lax.iota(dtype=dtype or jnp.result_type(start), size=size) * step + start
@@ -274,7 +267,7 @@ def concatenate(axis: AxisSelector, arrays: Sequence[NamedArray]) -> NamedArray:
     if axis_index is None:
         raise ValueError(f"Axis {aname} not found in 0th array {arrays[0]}")
 
-    axes: typing.Tuple[AxisSelector, ...] = arrays[0].axes
+    axes: tuple[AxisSelector, ...] = arrays[0].axes
     # we want to use the axis name for `axis`, because it's not uncommon for those to be different lengths in the arrays
     axes = axes[:axis_index] + (aname,) + axes[axis_index + 1 :]
     arrays = [a.rearrange(axes) for a in arrays]
@@ -931,9 +924,6 @@ __all__ = [
     "make_axes",
     "axis_name",
     "axis_size",
-    "NamedArrayAxesSpec",
-    "NamedArrayAxes",
-    "Named",
     "NamedArray",
     "broadcast_to",
     "broadcast_axis",
@@ -1106,4 +1096,14 @@ __all__ = [
     "is_named_array",
     "tree_checkpoint_name",
     "ScanCheckpointPolicy",
+    "quantization",
+    "util",
+    "einsum",
+    "broadcast_arrays",
+    "unflatten_axis",
+    "ReductionFunction",
+    "SimpleReductionFunction",
+    "NamedArrayAxes",
+    "NamedArrayAxesSpec",
+    "Named",
 ]

{haliax-1.4.dev364 → haliax-1.4.dev366}/src/haliax/_src/dot.py

@@ -4,7 +4,6 @@ import warnings
 from typing import Dict, Optional, Tuple
 
 import jax
-import jax.numpy as jnp
 
 import haliax
 from haliax.axis import (
@@ -20,7 +19,6 @@ from haliax.axis import (
 from haliax.core import NamedArray
 from haliax.jax_utils import _jittable_dg_einsum
 from haliax.types import DTypeLike, PrecisionLike
-from haliax.util import ensure_tuple
 
 
 # deprecated overload

{haliax-1.4.dev364 → haliax-1.4.dev366}/src/haliax/_src/fp8.py

@@ -1,4 +1,3 @@
-import warnings
 from functools import partial
 
 from jax import custom_jvp, custom_vjp, lax

{haliax-1.4.dev364 → haliax-1.4.dev366}/src/haliax/_src/state_dict.py

@@ -15,8 +15,7 @@ from jaxtyping import PyTree
 
 import haliax.partitioning as partitioning
 from haliax._src.util import index_where
-from haliax.axis import Axis
-from haliax.core import NamedArray, flatten_axes, named
+from haliax.core import NamedArray, named
 from haliax.jax_utils import is_jax_array_like, is_scalarish
 from haliax.tree_util import scan_aware_tree_map
 
@@ -209,14 +208,14 @@ def from_state_dict(tree: T, state_dict: StateDict, prefix: Optional[str] = None
             array = named(array, tree.axes)
             array = partitioning.auto_sharded(array)
 
-        return array
+        return array  # type: ignore
     elif is_jax_array_like(tree):
         if prefix is None:
             raise ValueError("Cannot extract a leaf value from a state dict without a prefix")
         # TODO: add "strict" flag so we can return None in cases where it's just missing
         return jnp.array(state_dict[prefix])
     elif tree is None:
-        return None
+        return None  # type: ignore
     else:
         if prefix is None:
             return tree

{haliax-1.4.dev364 → haliax-1.4.dev366}/src/haliax/core.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import contextlib
 import functools as ft
 import typing
@@ -32,7 +34,6 @@ from .axis import (
 )
 from .types import GatherScatterModeStr, IntScalar, PrecisionLike, Scalar
 
-
 NamedOrNumeric = Union[Scalar, "NamedArray"]
 NamedIndex = Union[int, slice_t, "NamedArray", dslice, list[int], jnp.ndarray]
 
@@ -181,19 +182,6 @@ def _parse_namedarray_axes(
     raise TypeError(f"Invalid NamedArray typing spec: {item}")
 
 
-class Named:
-    """Type annotation helper for :class:`NamedArray`.
-
-    ``Named["batch embed"]`` expands to ``Annotated[NamedArray, axes]`` so that
-    type checkers treat it as a ``NamedArray`` at static time while the axis
-    metadata is available at runtime via :func:`typing.get_args`.
-    """
-
-    def __class_getitem__(cls, item: NamedArrayAxesSpec) -> typing.Any:
-        axes = _parse_namedarray_axes(item)
-        return typing.Annotated[NamedArray, axes]
-
-
 class NamedArrayMeta(type):
     def __getitem__(cls, item: NamedArrayAxesSpec) -> typing.Any:
         axes = _parse_namedarray_axes(item)
@@ -2112,7 +2100,6 @@ def _convert_index_expr_to_dict(idx) -> dict[AxisSelector, NamedIndex]:
 __all__ = [
     "NamedArrayAxesSpec",
     "NamedArrayAxes",
-    "Named",
     "NamedArray",
     "named",
     "slice",

haliax-1.4.dev366/src/haliax/haxtyping.py

@@ -0,0 +1,138 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+
+@dataclass(frozen=True)
+class DTypeCategory:
+    """Represents a dtype category such as ``float`` or ``int``."""
+
+    name: str
+    category: Any
+
+    def __repr__(self) -> str:  # pragma: no cover - trivial
+        return self.name
+
+if TYPE_CHECKING:
+    # ── STATIC ONLY: re‑export jaxtyping’s aliases so mypy/Pyright/PyCharm see them
+    from jaxtyping import (
+        Float32 as f32, Float64 as f64, Float16 as f16, BFloat16 as bf16,
+        Int8 as i8, Int16 as i16, Int32 as i32, Int64 as i64,
+        UInt8 as u8, UInt16 as u16, UInt32 as u32, UInt64 as u64,
+        Bool as bool_, Complex64 as complex64, Complex128 as complex128,
+        Float as Float, Int as Int, UInt as UInt,
+    )
+    # axes‑only helper
+    from typing import Annotated as Named
+
+else:
+    # ── RUNTIME: custom wrappers for NamedArray, plus delegation to jaxtyping ──
+    import jaxtyping as jt
+    import jax.numpy as jnp
+    from typing import Annotated
+    from dataclasses import dataclass, replace
+
+    from .core import (
+        NamedArray,
+        NamedArrayAxes,
+        NamedArrayAxesSpec,
+        _parse_namedarray_axes,
+    )
+
+    def _with_dtype(axes: NamedArrayAxes, dtype):
+        """Attach dtype to axes metadata if not already set."""
+        return axes if axes.dtype is not None else replace(axes, dtype=dtype)
+
+    def _make_dtype_wrapper(dtype):
+        """Factory for f32, i32, etc."""
+        class _Wrapper:
+            def __class_getitem__(cls, item):
+                # two‑arg form: (BaseType, axes_spec)
+                if isinstance(item, tuple) and len(item) == 2:
+                    base, axes_spec = item
+                else:
+                    base, axes_spec = NamedArray, item
+
+                # Delegate non‑NamedArray to jaxtyping
+                if base is not NamedArray:
+                    # e.g. use jt.Float32 for jnp.float32
+                    jaxt = getattr(jt, f"Float{dtype.itemsize*8}") \
+                           if hasattr(dtype, "itemsize") else jt.Float
+                    return jaxt[base, axes_spec]
+
+                # Handle NamedArray path
+                axes = _parse_namedarray_axes(axes_spec)
+                return Annotated[NamedArray, _with_dtype(axes, dtype)]
+
+        return _Wrapper
+
+    # ── Build all dtype wrappers ─────────────────────────────────────────────
+    f32  = _make_dtype_wrapper(jnp.float32)
+    f64  = _make_dtype_wrapper(jnp.float64)
+    f16  = _make_dtype_wrapper(jnp.float16)
+    bf16 = _make_dtype_wrapper(jnp.bfloat16)
+
+    i8  = _make_dtype_wrapper(jnp.int8)
+    i16 = _make_dtype_wrapper(jnp.int16)
+    i32 = _make_dtype_wrapper(jnp.int32)
+    i64 = _make_dtype_wrapper(jnp.int64)
+
+    u8  = _make_dtype_wrapper(jnp.uint8)
+    u16 = _make_dtype_wrapper(jnp.uint16)
+    u32 = _make_dtype_wrapper(jnp.uint32)
+    u64 = _make_dtype_wrapper(jnp.uint64)
+
+    bool_      = _make_dtype_wrapper(jnp.bool_)
+    complex64  = _make_dtype_wrapper(jnp.complex64)
+    complex128 = _make_dtype_wrapper(jnp.complex128)
+
+
+
+
+    def _make_category_wrapper(name: str, category):
+        """Like _make_dtype_wrapper but matches any dtype in the JAX category."""
+
+        class _Wrapper:
+            def __class_getitem__(cls, item):
+                # same base/axes unpack logic
+                if isinstance(item, tuple) and len(item) == 2:
+                    base, axes_spec = item
+                else:
+                    base, axes_spec = NamedArray, item
+
+                # non‑NamedArray → delegate to jaxtyping’s category wrapper
+                if base is not NamedArray:
+                    return getattr(jt, name)[base, axes_spec]
+
+                # NamedArray path
+                axes = _parse_namedarray_axes(axes_spec)
+                cat = DTypeCategory(name, category)
+                return Annotated[NamedArray, _with_dtype(axes, cat)]
+
+        return _Wrapper
+
+
+    # Build the category wrappers
+    Float = _make_category_wrapper("float", jnp.floating)
+    Complex = _make_category_wrapper("complex", jnp.complexfloating)
+    Int = _make_category_wrapper("int", jnp.signedinteger)
+    UInt = _make_category_wrapper("uInt", jnp.unsignedinteger)
+
+
+    # ── Named: axes‑only helper ───────────────────────────────────────────────
+    class _NamedHelper:
+        @classmethod
+        def __class_getitem__(self, axes_spec_: tuple[type[NamedArray], NamedArrayAxesSpec]):
+            _, axes_spec = axes_spec_
+            axes = _parse_namedarray_axes(axes_spec)
+            return Annotated[NamedArray, axes]
+
+    Named = _NamedHelper
+
+    __all__ = [
+        "Named",
+        "f32", "f64", "f16", "bf16",
+        "i8", "i16", "i32", "i64",
+        "u8", "u16", "u32", "u64",
+        "bool_", "complex64", "complex128",
+    ]

{haliax-1.4.dev364 → haliax-1.4.dev366}/src/haliax/hof.py

@@ -23,7 +23,7 @@ from .axis import Axis, AxisSelection, AxisSelector, axis_spec_to_shape_dict, ax
 from .core import NamedArray
 from .jax_utils import Static, broadcast_prefix, is_jax_array_like
 from .partitioning import physical_axis_name
-from .util import ensure_tuple, is_named_array
+from .util import is_named_array
 
 
 def vmap(

{haliax-1.4.dev364 → haliax-1.4.dev366}/src/haliax/nn/attention.py

@@ -11,7 +11,6 @@ import haliax.random as hrandom
 from haliax.axis import Axis, AxisSelection, AxisSelector, AxisSpec, axis_name, axis_spec_to_shape_dict
 from haliax.core import NamedArray
 from haliax.types import PrecisionLike
-from haliax.util import ensure_tuple
 
 
 # With attention, we usually distinguish between the mask and the bias, though the former is just a special case of the

{haliax-1.4.dev364 → haliax-1.4.dev366}/src/haliax/nn/embedding.py

@@ -1,5 +1,4 @@
 import dataclasses
-import math
 import warnings
 from typing import Optional
 
@@ -12,7 +11,6 @@ from ..axis import Axis, AxisSpec, concat_axes
 from ..core import NamedArray
 from ..jax_utils import named_call
 from ..tree_util import resize_axis
-from ..util import ensure_tuple
 
 
 class Embedding(eqx.Module):

{haliax-1.4.dev364 → haliax-1.4.dev366}/src/haliax/nn/scan.py

@@ -2,7 +2,7 @@ import dataclasses
 import functools
 import re
 import warnings
-from typing import Any, Dict, Generic, Literal, Optional, Protocol, Sequence, Type, TypeVar, Union, cast
+from typing import Any, Dict, Generic, Optional, Protocol, Sequence, Type, TypeVar, cast
 
 import equinox as eqx
 import jax
@@ -169,7 +169,7 @@ class BlockSeq(ModuleWithStateDictSerialization, Generic[M]):
 
     @staticmethod
     def _slice_out(Block, i, x):
-        if haliax.is_named_array(x):
+        if isinstance(x, haliax.core.NamedArray):
             if haliax.selects_axis(x.axes, Block):
                 return x[Block, i]
             else:
@@ -411,7 +411,7 @@ class Stacked(ModuleWithStateDictSerialization, Generic[M]):
             else:
                 return tuple(x for _ in range(self.Block.size))
 
-        leaves, structure = jax.tree_util.tree_flatten(self.stacked, is_leaf=haliax.is_named_array)
+        leaves, structure = jax.tree_util.tree_flatten(self.stacked, is_leaf=haliax.util.is_named_array)
         unstacked_leaves = tuple(map(unbatch_leaf, leaves))
         # now we need to transpose the leaves
         unstacked_leaves = tuple(zip(*unstacked_leaves))

{haliax-1.4.dev364 → haliax-1.4.dev366}/src/haliax/partitioning.py

@@ -20,7 +20,7 @@ from .axis import Axis, AxisSelection, AxisSelector, axis_spec_to_shape_dict
 from .core import NamedArray
 from .jax_utils import Static, is_in_jit, is_jax_array_like, is_on_mac_metal
 from .tree_util import hashable_combine, hashable_partition
-from .util import StringHolderEnum, ensure_tuple
+from .util import StringHolderEnum
 
 
 PhysicalAxisSpec = Union[(str), Sequence[str]]

{haliax-1.4.dev364 → haliax-1.4.dev366}/src/haliax/random.py

@@ -1,7 +1,4 @@
 """Wrappers around jax.random functions."""
-import functools
-import inspect
-import warnings
 from typing import Optional
 
 import jax.random as jrandom