haliax 1.4.dev412__tar.gz → 1.4.dev419__tar.gz

{haliax-1.4.dev412 → haliax-1.4.dev419}/.agents/projects/api_parity.md

@@ -115,15 +115,15 @@ APIs that don't translate well to named tensors are intentionally omitted here.
 - [ ] `permute_dims`
 - [ ] `piecewise`
 - [ ] `place`
-- [ ] `poly`
-- [ ] `polyadd`
-- [ ] `polyder`
-- [ ] `polydiv`
-- [ ] `polyfit`
-- [ ] `polyint`
-- [ ] `polymul`
-- [ ] `polysub`
-- [ ] `polyval`
+- [x] `poly`
+- [x] `polyadd`
+- [x] `polyder`
+- [x] `polydiv`
+- [x] `polyfit`
+- [x] `polyint`
+- [x] `polymul`
+- [x] `polysub`
+- [x] `polyval`
 - [ ] `pow`
 - [ ] `promote_types`
 - [ ] `put`
@@ -134,7 +134,7 @@ APIs that don't translate well to named tensors are intentionally omitted here.
 - [ ] `resize`
 - [ ] `result_type`
 - [ ] `rollaxis`
-- [ ] `roots`
+- [x] `roots`
 - [ ] `rot90`
 - [ ] `select`
 - [ ] `setdiff1d`
@@ -148,7 +148,7 @@ APIs that don't translate well to named tensors are intentionally omitted here.
 - [ ] `tri`
 - [ ] `tril_indices`
 - [ ] `tril_indices_from`
-- [ ] `trim_zeros`
+- [x] `trim_zeros`
 - [ ] `triu_indices`
 - [ ] `triu_indices_from`
 - [ ] `union1d`
@@ -156,7 +156,7 @@ APIs that don't translate well to named tensors are intentionally omitted here.
 - [ ] `unravel_index`
 - [ ] `unstack`
 - [ ] `unwrap`
-- [ ] `vander`
+- [x] `vander`
 - [ ] `vsplit`
 - [ ] `vstack`
 

{haliax-1.4.dev412 → haliax-1.4.dev419}/AUTHORS.md

@@ -3,3 +3,4 @@
 The Levanter Authors currently include:
 
 - The Board of Trustees of the Leland Stanford Junior University
+- Open Athena AI Foundation, Inc.

{haliax-1.4.dev412 → haliax-1.4.dev419}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev412
+Version: 1.4.dev419
 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/
@@ -17,6 +17,7 @@ Classifier: Programming Language :: Python :: 3
 Requires-Python: >=3.10
 Requires-Dist: aqtp>=0.8.2
 Requires-Dist: equinox>=0.10.6
+Requires-Dist: jax>=0.6.2
 Requires-Dist: jaxtyping>=0.2.20
 Requires-Dist: jmp>=0.0.4
 Requires-Dist: safetensors>=0.4.3

{haliax-1.4.dev412 → haliax-1.4.dev419}/docs/api.md

@@ -6,19 +6,89 @@ arrays (see [haliax.zeros][] and [haliax.ones][]) as well as things like reducti
 
 ## Axis Types
 
-There are four types related to [haliax.Axis][] you will see in the API reference:
+If you already speak NumPy or `jax.numpy`, think of Haliax as swapping positional axes (`axis=0`) for named axes
+(`axis="batch"`). The type hints in this section describe the different ways those named axes can be provided to the API.
+They appear throughout the documentation and in signatures so that you can quickly tell which forms are accepted.
 
-* [haliax.Axis][]: This is the main type for representing axes. It is a dataclass with a name and a size.
-* [haliax.AxisSelector][]: This is a type alias for either [haliax.Axis][] or a `str`. This type is used when we want
-  one axis and the size can be inferred from the inputs.
-* [haliax.AxisSpec][]: This is a type alias for either [haliax.Axis][] or a tuple/list of [haliax.Axis][]. This type is
-  used when we want one or more axes and the sizes cannot be inferred from the inputs, for instance when creating arrays.
-* [haliax.AxisSelection][]: This is a type alias for either [haliax.AxisSelector][] or a tuple/list of [haliax.AxisSelector][].
-    This type is used when we want one or more axes and the sizes can be inferred from the inputs, for instance when
-    reducing an array.
+| Name | Accepts | When to use it | Example |
+| --- | --- | --- | --- |
+| [`Axis`][haliax.Axis] | `Axis(name: str, size: int)` | Define a named dimension with an explicit size | `Batch = Axis("batch", 32)` |
+| [`AxisSelector`][haliax.AxisSelector] | `Axis` or `str` | Refer to an existing axis whose size can be inferred from the arrays you pass in | `x.sum(axis="batch")` |
+| [`AxisSpec`][haliax.AxisSpec] | `dict[str, int]`, `Axis`, or a sequence of `Axis` objects | Create or reshape arrays when the axis sizes must be provided | `hax.zeros((Batch, Feature))` |
+| [`AxisSelection`][haliax.AxisSelection] | `dict[str, int | None]`, `AxisSpec`, or a sequence of `AxisSelector` values | Work with one or more existing axes (reductions, indexing helpers, flattening, …) | `x.sum(axis=("batch", Feature))` |
 
-Occasionally, an axis size can be inferred in some circumstances but not others. When this happens, we still use
-`AxisSelector` but document the behavior in the docstring. A RuntimeError will be raised if the size cannot be inferred.
+### Axis
+
+An [`Axis`][haliax.Axis] is the fundamental building block: it is a tiny dataclass that stores a name and a size.  You can
+construct one directly or use [`haliax.make_axes`][] to generate several at a time.
+
+```python
+import haliax as hax
+from haliax import Axis
+
+Batch = Axis("batch", 32)
+Feature = Axis("feature", 128)
+x = hax.ones((Batch, Feature))
+print(Batch.name, Batch.size)
+```
+
+Using `Axis` objects keeps array creation explicit and gives reusable handles you can share between different tensors.
+Equality compares both the name and size so you get guardrails when wiring pieces together.
+
+### AxisSelector
+
+An [`AxisSelector`][haliax.AxisSelector] accepts either an `Axis` object or just the axis name as a string.  It is used
+whenever a function can read the axis size from one of its arguments.  This mirrors how NumPy lets you pass `axis=0` when
+reducing an array:
+
+```python
+total = x.sum(axis=Batch)       # using the Axis handle
+same_total = x.sum(axis="batch")  # using only the name
+```
+
+Strings are convenient when you only care about the name, but `Axis` objects still work so you can keep using the handles
+you created earlier.  If an axis with that name is missing, Haliax raises a `ValueError`.
+
+### AxisSpec
+
+An [`AxisSpec`][haliax.AxisSpec] is used when Haliax needs full size information to create or reshape an array.  You can
+provide a shape dictionary (sometimes called a "shape dict") that maps names to sizes, or a sequence of `Axis` objects:
+
+```python
+shape = {"batch": 32, "feature": 128}
+y = hax.zeros(shape)                 # using a shape dict
+z = hax.zeros((Batch, Feature))      # using the Axis objects directly
+```
+
+Both forms describe the same layout.  Python dictionaries preserve insertion order, so the ordering in a shape dict matches
+the order that axes appear in the array.  Sequences must contain `Axis` objects (not plain strings) because Haliax cannot
+otherwise know the axis sizes.
+
+### AxisSelection
+
+[`AxisSelection`][haliax.AxisSelection] generalizes the previous aliases so you can talk about several axes at once.  It
+shows up in reductions, indexing helpers, axis-mapping utilities, and anywhere you might have written `axis=(0, 1)` in
+NumPy.  You may supply:
+
+* a sequence mixing `Axis` objects and strings, e.g. `("batch", Feature)` when reducing two axes,
+* an `AxisSpec`, which is handy when you already have a tuple of `Axis` objects, or
+* a "partial shape dict" where the values are either sizes or `None` to indicate "any size".  Dictionaries are useful when
+  you only care about a subset of axes or want to assert a particular size.
+
+```python
+# Reduce over two axes using a tuple of selectors.
+scalar = x.sum(axis=("batch", Feature))
+
+# Ask for the axes by name and optionally pin sizes.
+x.resolve_axis({"batch": None, "feature": None})  # returns {"batch": 32, "feature": 128}
+
+from haliax.axis import selects_axis
+assert selects_axis((Batch, "feature"), {"batch": None, "feature": 128})
+```
+
+Occasionally, an axis size can be inferred in some circumstances but not others. When this happens we still use
+`AxisSelector` (or `AxisSelection` for multiple axes) but document the behavior in the docstring. A `RuntimeError` will be
+raised if the size cannot be inferred.
 
 ::: haliax.Axis
 ::: haliax.AxisSelector
@@ -270,6 +340,21 @@ These are all more or less directly from JAX's NumPy API.
 ::: haliax.subtract
 ::: haliax.true_divide
 
+### Polynomial Operations
+
+::: haliax.poly
+::: haliax.polyadd
+::: haliax.polysub
+::: haliax.polymul
+::: haliax.polydiv
+::: haliax.polyint
+::: haliax.polyder
+::: haliax.polyval
+::: haliax.polyfit
+::: haliax.roots
+::: haliax.trim_zeros
+::: haliax.vander
+
 ### Other Operations
 
 ::: haliax.bincount

haliax-1.4.dev419/docs/primer.md

@@ -0,0 +1,183 @@
+# Haliax Primer
+
+Haliax provides named tensors built on top of JAX.  This primer is written for LLM agents and other downstream libraries and collects the core ideas for quick reference.
+
+## Axes and Named Arrays
+
+Arrays are indexed by `Axis` objects. You can define them explicitly or generate several with `make_axes`.
+You may also specify shapes with a **shape dict**, mapping axis names to sizes.
+
+```python
+import haliax as hax
+from haliax import Axis
+
+Batch = Axis("batch", 4)
+Feature = Axis("feature", 8)
+# or: Batch, Feature = hax.make_axes(batch=4, feature=8)
+# using Axis objects
+x = hax.zeros((Batch, Feature))
+# or using a shape dict
+shape = {"batch": 4, "feature": 8}
+x = hax.zeros(shape)
+```
+
+Most functions accept either axes or shape dicts interchangeably.
+
+A tensor with named axes is a [`NamedArray`][haliax.NamedArray]. Elementwise operations mirror `jax.numpy` but accept named axes.
+
+### Ways to Describe Axes
+
+If you're used to `axis=0` style code in NumPy or JAX, think of Haliax as swapping those positional indices for names
+like `axis="batch"`. The API hints refer to a few helper aliases; this table shows how they map back to familiar
+concepts:
+
+| Name | Accepts | Typical use | Example |
+| --- | --- | --- | --- |
+| [`Axis`][haliax.Axis] | `Axis(name: str, size: int)` | Define a named dimension with a fixed size | `Batch = Axis("batch", 32)` |
+| [`AxisSelector`][haliax.AxisSelector] | `Axis` or `str` | Refer to an existing axis when the size can be inferred from the argument | `x.sum(axis="batch")` |
+| [`AxisSpec`][haliax.AxisSpec] | `dict[str, int]`, `Axis`, or a sequence of `Axis` objects | Supply complete shape information (array creation, reshaping) | `hax.zeros((Batch, Feature))` |
+| [`AxisSelection`][haliax.AxisSelection] | `dict[str, int | None]`, `AxisSpec`, or a sequence of `AxisSelector` values | Work with one or more existing axes (reductions, indexing helpers, flattening, …) | `x.sum(axis=("batch", Feature))` |
+
+The following sections expand on each alias with quick references and NumPy-style parallels.
+
+#### `Axis`: reusable named dimensions
+
+An [`Axis`][haliax.Axis] stores a `name` and a `size`. Create them directly or let
+[`haliax.make_axes`][] build a handful at once. Because axes compare by both name and size, they act as reusable handles and
+catch many wiring mistakes early.
+
+```python
+Batch = Axis("batch", 32)
+Feature = Axis("feature", 128)
+x = hax.ones((Batch, Feature))
+print(Batch.name, Batch.size)
+```
+
+#### `AxisSelector`: when the size is already known
+
+Many functions already see the array whose axes you're referencing (e.g. reductions). In those cases you can pass either the
+`Axis` object or simply the axis name as a string. Haliax resolves the name against the array, similar to `axis=0` in NumPy.
+
+```python
+total = x.sum(axis=Batch)       # use the Axis handle
+same_total = x.sum(axis="batch")  # or just the name
+```
+
+If you reference an axis name that isn't present, Haliax raises a `ValueError`.
+
+#### `AxisSpec`: describing complete shapes
+
+When Haliax needs explicit sizes—creating arrays, reshaping, broadcasting to a new axis—you provide an [`AxisSpec`][haliax.AxisSpec].
+Shape dictionaries keep things close to standard Python, while sequences require actual `Axis` objects so the sizes stay explicit.
+
+```python
+shape = {"batch": 32, "feature": 128}
+y = hax.zeros(shape)                 # using a shape dict
+z = hax.zeros((Batch, Feature))      # or a sequence of Axis objects
+```
+
+Python dictionaries preserve insertion order, so the layout in a shape dict matches the order of axes in the resulting array.
+
+#### `AxisSelection`: several axes at once
+
+[`AxisSelection`][haliax.AxisSelection] is the plural form used by reductions, indexing helpers, and flattening utilities. Supply a
+tuple mixing `Axis` objects and strings, reuse an existing `AxisSpec`, or pass a partial shape dict where values are either
+sizes or `None` for "any size".
+
+```python
+scalar = x.sum(axis=("batch", Feature))
+
+# Ask for the axes by name and optionally pin sizes.
+x.resolve_axis({"batch": None, "feature": None})  # returns {"batch": 32, "feature": 128}
+```
+
+Partial shape dicts shine when you only care about a subset of axes or want assertions about their sizes. If the axis size
+cannot be inferred from the provided arguments, Haliax raises a `RuntimeError`.
+
+## Indexing and Broadcasting
+
+Use axis names when slicing. Dictionaries are convenient for several axes:
+
+```python
+first = x["batch", 0]
+sub = x["batch", 1:3]
+# or with a dict
+first = x[{"batch": 0}]
+sub = x[{"batch": slice(1, 3)}]
+```
+
+Axes broadcast by matching names. `broadcast_axis` adds a new axis to an array:
+
+```python
+row = hax.arange(Feature)
+outer = row.broadcast_axis(Batch) * hax.arange(Batch)
+```
+
+See [Indexing and Slicing](indexing.md) and [Broadcasting](broadcasting.md) for details.
+
+## Rearranging Axes
+
+`rearrange` changes axis order and can merge or split axes using einops‑style syntax.  It is useful when interfacing with positional APIs.
+
+```python
+# transpose features and batch
+x_t = hax.rearrange(x, "batch feature -> feature batch")
+```
+
+More examples appear in [Rearrange](rearrange.md).
+
+## Matrix Multiplication
+
+`dot` contracts over named axes while preserving order independence.
+
+```python
+Weight = Axis("weight", 8)
+w = hax.ones((Feature, Weight))
+prod = hax.dot(x, w, axis=Feature)
+```
+
+For more complex contractions use [`einsum`][haliax.einsum].  See [Matrix Multiplication](matmul.md).
+
+## Scans and Folds
+
+Use [`scan`][haliax.scan] or [`fold`][haliax.fold] to apply a function along an axis with optional gradient checkpointing.
+
+```python
+Time = Axis("time", 10)
+sequence = hax.ones((Time, Feature))
+
+def add(prev, cur):
+    return prev + cur
+
+result = hax.fold(add, Time)(hax.zeros((Feature,)), sequence)
+```
+
+See [Scan and Fold](scan.md) for checkpointing policies and stacked modules.
+
+## Partitioning
+
+Arrays and modules can be distributed across devices by mapping named axes to mesh axes:
+
+```python
+with hax.axis_mapping({"batch": "data"}):
+    sharded = hax.shard(x)
+```
+
+The [Partitioning](partitioning.md) guide explains how to set up device meshes and shard arrays.
+
+## Typing Support
+
+Type annotations use `haliax.haxtyping` which extends `jaxtyping`:
+
+```python
+import haliax.haxtyping as ht
+
+def f(t: ht.Float[hax.NamedArray, "batch feature"]):
+    ...
+```
+
+See [Typing](typing.md) for matching runtime checks and dtype-aware annotations.
+
+---
+
+This primer highlights common patterns.  The [cheatsheet](cheatsheet.md) lists many additional conversions from JAX to Haliax.

{haliax-1.4.dev412 → haliax-1.4.dev419}/pyproject.toml

@@ -21,8 +21,7 @@ classifiers = [
     "Intended Audience :: Science/Research",
 ]
 dependencies = [
-    # we require that you install jax yourself, since the extras vary by system.
-    # jax = {version = ">=0.4.19,<0.5.0"}
+     "jax >= 0.6.2",
     "equinox>=0.10.6",
     "jaxtyping>=0.2.20",
     "jmp>=0.0.4",

{haliax-1.4.dev412 → haliax-1.4.dev419}/src/haliax/__about__.py

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

{haliax-1.4.dev412 → haliax-1.4.dev419}/src/haliax/__init__.py

@@ -8,11 +8,8 @@ from typing import Optional, Sequence
 
 import jax
 import jax.numpy as jnp
+from jax.typing import DTypeLike
 
-try:
-    from jax.typing import DTypeLike
-except ImportError:
-    from jax._src.typing import DTypeLike
 
 import haliax.debug as debug
 import haliax.nn as nn
@@ -96,6 +93,22 @@ from .ops import (
     bincount,
     where,
 )
+
+from .poly import (
+    poly,
+    polyadd,
+    polysub,
+    polymul,
+    polydiv,
+    polyint,
+    polyder,
+    polyval,
+    polyfit,
+    roots,
+    trim_zeros,
+    vander,
+)
+
 from .fft import (
     fft,
     fftfreq,
@@ -108,7 +121,7 @@ from .fft import (
     rfft,
     rfftfreq,
 )
-from .partitioning import auto_sharded, axis_mapping, fsdp, named_jit, shard, shard_with_axis_mapping
+from .partitioning import auto_sharded, axis_mapping, fsdp, named_jit, set_mesh, shard, shard_with_axis_mapping
 from .specialized_fns import top_k
 from .types import Scalar
 from .util import is_named_array
@@ -715,7 +728,9 @@ def nanmean(
     where: Optional[NamedArray] = None,
     dtype: Optional[DTypeLike] = None,
 ) -> NamedArray:
-    return wrap_reduction_call(jnp.nanmean, array, axis, where, single_axis_only=False, supports_where=True, dtype=dtype)
+    return wrap_reduction_call(
+        jnp.nanmean, array, axis, where, single_axis_only=False, supports_where=True, dtype=dtype
+    )
 
 
 def nanmin(
@@ -734,7 +749,9 @@ def nanprod(
     where: Optional[NamedArray] = None,
     dtype: Optional[DTypeLike] = None,
 ) -> NamedArray:
-    return wrap_reduction_call(jnp.nanprod, array, axis, where, single_axis_only=False, supports_where=True, dtype=dtype)
+    return wrap_reduction_call(
+        jnp.nanprod, array, axis, where, single_axis_only=False, supports_where=True, dtype=dtype
+    )
 
 
 def nanstd(
@@ -745,7 +762,9 @@ def nanstd(
     ddof: int = 0,
     dtype: Optional[DTypeLike] = None,
 ) -> NamedArray:
-    return wrap_reduction_call(jnp.nanstd, array, axis, where, single_axis_only=False, supports_where=True, dtype=dtype, ddof=ddof)
+    return wrap_reduction_call(
+        jnp.nanstd, array, axis, where, single_axis_only=False, supports_where=True, dtype=dtype, ddof=ddof
+    )
 
 
 def nansum(
@@ -755,7 +774,9 @@ def nansum(
     where: Optional[NamedArray] = None,
     dtype: Optional[DTypeLike] = None,
 ) -> NamedArray:
-    return wrap_reduction_call(jnp.nansum, array, axis, where, single_axis_only=False, supports_where=True, dtype=dtype)
+    return wrap_reduction_call(
+        jnp.nansum, array, axis, where, single_axis_only=False, supports_where=True, dtype=dtype
+    )
 
 
 def nanvar(
@@ -766,7 +787,9 @@ def nanvar(
     ddof: int = 0,
     dtype: Optional[DTypeLike] = None,
 ) -> NamedArray:
-    return wrap_reduction_call(jnp.nanvar, array, axis, where, single_axis_only=False, supports_where=True, dtype=dtype, ddof=ddof)
+    return wrap_reduction_call(
+        jnp.nanvar, array, axis, where, single_axis_only=False, supports_where=True, dtype=dtype, ddof=ddof
+    )
 
 
 # "Normalization" functions that use an axis but don't change the shape
@@ -807,14 +830,17 @@ def sort(a: NamedArray, axis: AxisSelector) -> NamedArray:
     return wrap_axiswise_call(jnp.sort, a, axis, single_axis_only=True)
 
 
-def argsort(a: NamedArray, axis: AxisSelector) -> NamedArray:
+def argsort(a: NamedArray, axis: AxisSelector | None, *, stable: bool = False) -> NamedArray:
     """
     Named version of [jax.numpy.argsort](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.argsort.html).
 
     If `axis` is None, the returned array will be a 1D array of indices that would sort the flattened array,
     identical to `jax.numpy.argsort(a.array)`.
+
+    Args:
+        stable: If ``True``, ensures that the indices of equal elements preserve their relative order.
     """
-    return wrap_axiswise_call(jnp.argsort, a, axis, single_axis_only=True)
+    return wrap_axiswise_call(jnp.argsort, a, axis, single_axis_only=True, stable=stable)
 
 
 # elemwise binary ops
@@ -1226,6 +1252,18 @@ __all__ = [
     "clip",
     "tril",
     "triu",
+    "poly",
+    "polyadd",
+    "polysub",
+    "polymul",
+    "polydiv",
+    "polyint",
+    "polyder",
+    "polyval",
+    "polyfit",
+    "roots",
+    "trim_zeros",
+    "vander",
     "fft",
     "ifft",
     "hfft",
@@ -1311,4 +1349,5 @@ __all__ = [
     "NamedArrayAxes",
     "NamedArrayAxesSpec",
     "Named",
+    "set_mesh",
 ]

{haliax-1.4.dev412 → haliax-1.4.dev419}/src/haliax/_src/state_dict.py

@@ -197,7 +197,7 @@ def from_state_dict(tree: T, state_dict: StateDict, prefix: Optional[str] = None
         array = state_dict[prefix]
 
         if isinstance(array, np.ndarray):
-            mesh = partitioning._get_mesh()
+            mesh = jax.sharding.get_abstract_mesh()
             # TODO: modernize this
             if jax.device_count() > 1:  # this happens with the default mesh
                 pspec = partitioning.pspec_for_axis(tree.axes)

{haliax-1.4.dev412 → haliax-1.4.dev419}/src/haliax/core.py

@@ -623,8 +623,8 @@ class NamedArray(metaclass=NamedArrayMeta):
     def argmin(self, axis: Optional[AxisSelector]) -> "NamedArray":  # pragma: no cover
         return haliax.argmin(self, axis=axis)
 
-    def argsort(self, axis: AxisSelector) -> "NamedArray":  # pragma: no cover
-        return haliax.argsort(self, axis=axis)
+    def argsort(self, axis: AxisSelector | None, *, stable: bool = False) -> "NamedArray":  # pragma: no cover
+        return haliax.argsort(self, axis=axis, stable=stable)
 
     def astype(self, dtype) -> "NamedArray":  # pragma: no cover
         return NamedArray(self.array.astype(dtype), self.axes)
@@ -1183,6 +1183,26 @@ def index(array: NamedArray, slices: Mapping[AxisSelector, NamedIndex]) -> Named
 def _compute_new_axes_and_slices_for_index(
     array, slices
 ) -> tuple[AxisSpec, list[py_slice | dslice | jnp.ndarray | int | list[int]]]:
+    def _is_integer_like_scalar_index(value: Any) -> bool:
+        if isinstance(value, (int, np.integer)):
+            return True
+        if not is_jax_array_like(value):
+            return False
+        shape = getattr(value, "shape", None)
+        if shape != ():
+            return False
+        dtype = getattr(value, "dtype", None)
+        if dtype is None:
+            return False
+        return jnp.issubdtype(dtype, jnp.integer)
+
+    def _coerce_integer_like_index(value: Any):
+        if isinstance(value, np.integer):
+            return int(value)
+        if type(value) is np.ndarray and value.shape == () and jnp.issubdtype(value.dtype, jnp.integer):
+            return int(value.item())
+        return value
+
     ordered_slices: list = [py_slice(None, None, None)] * len(array.axes)  # type: ignore
     kept_axes = [True] * len(array.axes)
     array_slice_indices = []
@@ -1205,8 +1225,8 @@ def _compute_new_axes_and_slices_for_index(
             kept_axes[axis_index] = False
             array_slice_indices.append(axis_index)
             index_axis_names.add(orig_axis.name)
-        elif isinstance(slice_, int):
-            ordered_slices[axis_index] = slice_
+        elif _is_integer_like_scalar_index(slice_):
+            ordered_slices[axis_index] = _coerce_integer_like_index(slice_)
             kept_axes[axis_index] = False
         elif isinstance(slice_, NamedArray):
             ordered_slices[axis_index] = slice_

{haliax-1.4.dev412 → haliax-1.4.dev419}/src/haliax/fft.py

@@ -1,3 +1,7 @@
+# Copyright 2025 The Levanter Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
 """Named wrappers around :mod:`jax.numpy.fft`.
 
 These functions mirror the behaviour of their :mod:`jax.numpy.fft` counterparts
@@ -43,7 +47,6 @@ import jax.numpy.fft as jfft
 from .axis import Axis, AxisSelector, AxisSelection
 from .core import NamedArray
 
-
 AxisSizeLike = int | Axis | None
 AxisMapping = Mapping[AxisSelector, AxisSizeLike]
 

{haliax-1.4.dev412 → haliax-1.4.dev419}/src/haliax/nn/__init__.py

@@ -3,6 +3,8 @@
 # SPDX-License-Identifier: Apache-2.0
 
 
+import numbers
+
 import jax.nn as jnn
 import jax.numpy as jnp
 
@@ -64,16 +66,25 @@ def one_hot(x: NamedArray | int, class_axis: Axis, *, dtype=None) -> NamedArray:
         # Disabling this to prevent a crash in XLA on GPU
         # return hax.auto_sharded(hax.named(array, x.axes + (class_axis,)))
         return hax.named(array, x.axes + (class_axis,))
-    else:
-        assert isinstance(x, int)
+
+    x_array = jnp.asarray(x)
+    if x_array.ndim != 0:
+        raise TypeError("one_hot expects a scalar integer or NamedArray, " f"but got array with shape {x_array.shape}")
+
+    if isinstance(x, numbers.Integral):
         assert class_axis.size > x >= -class_axis.size
+    else:
+        if not (jnp.issubdtype(x_array.dtype, jnp.integer) or jnp.issubdtype(x_array.dtype, jnp.bool_)):
+            raise TypeError(
+                "one_hot expects an integer scalar or NamedArray, " f"but got value with dtype {x_array.dtype}"
+            )
 
-        one = 1
-        if dtype is not None:
-            one = dtype(one)
+    one = 1
+    if dtype is not None:
+        one = dtype(one)
 
-        array = jnp.zeros(class_axis.size, dtype=dtype).at[x].set(one)
-        return hax.auto_sharded(haliax.named(array, class_axis))
+    array = jnp.zeros(class_axis.size, dtype=dtype).at[x_array].set(one)
+    return hax.auto_sharded(haliax.named(array, class_axis))
 
 
 __all__ = [

{haliax-1.4.dev412 → haliax-1.4.dev419}/src/haliax/nn/linear.py

@@ -251,7 +251,7 @@ def _gmm(lhs, rhs, group_sizes, out_axes, sharded=False, ar=False):
     else:
         gmm_fn = shard_map(
             partial(gmm_sharded, ar=ar),
-            mesh=hax.partitioning._get_mesh(),
+            mesh=jax.sharding.get_abstract_mesh(),
             in_specs=(
                 hax.partitioning.pspec_for_axis(lhs.axes),
                 hax.partitioning.pspec_for_axis(rhs.axes),