haliax 1.4.dev413__tar.gz → 1.4.dev419__tar.gz

{haliax-1.4.dev413 → 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.dev413 → 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.dev413 → haliax-1.4.dev419}/PKG-INFO

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

@@ -340,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.dev413 → 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.dev413 → haliax-1.4.dev419}/src/haliax/__about__.py

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

{haliax-1.4.dev413 → 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.dev413 → 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.dev413 → 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.dev413 → 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.dev413 → 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.dev413 → 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),

{haliax-1.4.dev413 → haliax-1.4.dev419}/src/haliax/ops.py

@@ -127,9 +127,7 @@ def nonzero(array: NamedArray, *, size: Axis, fill_value: int = 0) -> tuple[Name
     if not isinstance(array, NamedArray):
         raise ValueError("array must be a NamedArray")
 
-    return tuple(
-        NamedArray(idx, (size,)) for idx in jnp.nonzero(array.array, size=size.size, fill_value=fill_value)
-    )
+    return tuple(NamedArray(idx, (size,)) for idx in jnp.nonzero(array.array, size=size.size, fill_value=fill_value))
 
 
 def clip(array: NamedOrNumeric, a_min: NamedOrNumeric, a_max: NamedOrNumeric) -> NamedArray:
@@ -335,7 +333,7 @@ def unique(
     )
 
     if axis is not None:
-        axis_index = array._lookup_indices(axis)
+        axis_index = array.axis_indices(axis)
         if axis_index is None:
             raise ValueError(f"Axis {axis} not found in array. Available axes: {array.axes}")
         out = jnp.unique(array.array, axis=axis_index, **kwargs)

{haliax-1.4.dev413 → haliax-1.4.dev419}/src/haliax/partitioning.py

@@ -10,31 +10,13 @@ import threading
 import typing
 import warnings
 from math import prod
-from typing import Callable, ContextManager, Mapping, Optional, ParamSpec, Sequence, TypeVar, Union
+from typing import Callable, ContextManager, Mapping, Optional, ParamSpec, Sequence, TypeVar, Union, cast
 
 import equinox as eqx
 import jax
 from equinox import is_array, module_update_wrapper
 from jax.lax import with_sharding_constraint
-from jax.sharding import (
-    Mesh,
-    NamedSharding,
-    PartitionSpec,
-    SingleDeviceSharding,
-)
-
-try:  # jax>=0.4.26
-    from jax.sharding import AbstractMesh, get_abstract_mesh
-except Exception:  # pragma: no cover - older JAX versions
-    AbstractMesh = Mesh  # type: ignore[misc,assignment]
-
-    def get_abstract_mesh():  # type: ignore[dead-code]
-        try:
-            from jax.interpreters.pxla import thread_resources
-        except Exception:
-            from jax.experimental.maps import thread_resources
-
-        return thread_resources.env.physical_mesh
+from jax.sharding import AbstractMesh, NamedSharding, Mesh, PartitionSpec, SingleDeviceSharding, get_abstract_mesh
 
 
 from jaxtyping import PyTree
@@ -50,6 +32,7 @@ from .util import StringHolderEnum
 
 PhysicalAxisSpec = Union[(str), Sequence[str]]
 ResourceMapping = Mapping[(str), PhysicalAxisSpec]
+MeshLike = Union[Mesh, AbstractMesh]
 """Mapping from logical axis names to physical axis names"""
 
 F = typing.TypeVar("F", bound=typing.Callable)
@@ -109,6 +92,71 @@ def current_thread_local_mapping():
     return _mapping_holder.thread_data.resource_mapping
 
 
+def _resolve_mesh(mesh: Optional[MeshLike] = None) -> Optional[MeshLike]:
+    """Inside jit, prefer an abstract mesh, outside jit prefer a concrete mesh."""
+
+    from jax._src.mesh import get_concrete_mesh
+
+    if mesh is not None:
+        if is_in_jit() and isinstance(mesh, Mesh):
+            return mesh.abstract_mesh
+        return mesh
+
+    if is_in_jit():
+        abstract = get_abstract_mesh()
+        if not abstract or abstract.empty:
+            concrete = get_concrete_mesh()
+            if concrete is not None and not concrete.empty:
+                return concrete.abstract_mesh
+
+        from jax.interpreters.pxla import thread_resources
+
+        old_mesh = thread_resources.env.physical_mesh
+        if old_mesh is not None and not old_mesh.empty:
+            return old_mesh.abstract_mesh
+
+        return abstract
+    else:
+        mesh = get_concrete_mesh() or get_abstract_mesh()
+        if mesh is not None and not mesh.empty:
+            return mesh
+
+        from jax.interpreters.pxla import thread_resources
+
+        old_mesh = thread_resources.env.physical_mesh
+        if old_mesh is not None and not old_mesh.empty:
+            return old_mesh
+
+    return None
+
+
+def mesh_context(mesh: MeshLike) -> ContextManager[None]:
+    """Context manager that normalizes mesh handling across JAX versions."""
+
+    set_mesh_fn = getattr(jax, "set_mesh", None)
+    use_mesh_fn = getattr(jax.sharding, "use_mesh", None)
+
+    manager_factory: Optional[Callable[[MeshLike], ContextManager[None]]] = None
+    if set_mesh_fn is not None:
+        manager_factory = cast(Callable[[MeshLike], ContextManager[None]], set_mesh_fn)
+    elif use_mesh_fn is not None:
+        manager_factory = cast(Callable[[MeshLike], ContextManager[None]], use_mesh_fn)
+
+    if manager_factory is None:
+        msg = "Haliax requires a version of JAX that provides either `jax.set_mesh` or `jax.sharding.use_mesh`."
+        raise RuntimeError(msg)
+
+    context_manager = manager_factory(mesh)
+
+    return context_manager
+
+
+def set_mesh(mesh: MeshLike) -> ContextManager[None]:
+    """Compatibility wrapper around `mesh_context` matching the JAX 0.7 API."""
+
+    return mesh_context(mesh)
+
+
 def auto_sharded(x: T, mesh: Optional[Mesh] = None) -> T:
     """
     Shard a PyTree using the global axis mapping. NamedArrays in the PyTree are sharded using the axis mapping
@@ -143,11 +191,15 @@ def shard(x: T, mapping: Optional[ResourceMapping] = None, mesh: Optional[Mesh]
 
     assert not isinstance(mesh, dict)
 
-    if mesh is None:
-        mesh = _get_mesh()
+    resolved_mesh = _resolve_mesh(mesh)
 
-        if mesh.empty:
-            return x
+    if resolved_mesh is None:
+        if not is_in_jit():
+            warnings.warn("No mesh found. Not sharding.", RuntimeWarning)
+        return x
+
+    if isinstance(resolved_mesh, AbstractMesh) and resolved_mesh.empty:
+        return x
 
     if is_in_jit() and is_on_mac_metal():
         warnings.warn("Sharding constraints are not supported in jit on metal", RuntimeWarning)
@@ -162,16 +214,11 @@ def shard(x: T, mapping: Optional[ResourceMapping] = None, mesh: Optional[Mesh]
             # could use eqx.partition to avoid this, but eh
             return named
 
-        sharding = infer_resource_partitions(named, mapping, mesh=mesh, preserve_existing_shardings=False)
-        assert isinstance(sharding, NamedSharding)
-        in_sharding = getattr(named.array, "sharding", None)
+        pspec = pspec_for(named, mapping, preserve_existing_shardings=False)
+        assert isinstance(pspec, PartitionSpec)
+        sharding = NamedSharding(resolved_mesh, pspec)
         if is_in_jit():
             return with_sharding_constraint(named, sharding)
-        # as a special case, SingleDeviceShardings are routed through jit
-        elif isinstance(in_sharding, SingleDeviceSharding) and in_sharding._device in sharding.device_set:
-            # TODO(dlwh): this should be unnecessary in JAX soon. Check after 2024-08-01
-            sharded_array = jax.jit(lambda x: x, out_shardings=sharding)(named)
-            return sharded_array
         else:
             ret = jax.device_put(named, sharding)
             return ret
@@ -189,7 +236,6 @@ def pspec_for(
     tree: PyTree,
     resource_mapping: Optional[ResourceMapping] = None,
     preserve_existing_shardings: bool = True,
-    use_auto_sharding: bool = True,
 ) -> PyTree:
     """Infer the :class:`PartitionSpec` for a module.
 
@@ -276,7 +322,6 @@ def infer_resource_partitions(
     tree: PyTree,
     resource_mapping: Optional[ResourceMapping] = None,
     preserve_existing_shardings: bool = True,
-    use_auto_sharding: bool = True,
     mesh: Optional[Mesh] = None,
 ) -> PyTree:
     """
@@ -291,11 +336,12 @@ def infer_resource_partitions(
         tree,
         resource_mapping=resource_mapping,
         preserve_existing_shardings=preserve_existing_shardings,
-        use_auto_sharding=use_auto_sharding,
     )
 
-    mesh = mesh or _get_mesh()
-    assert not isinstance(mesh, dict)
+    resolved_mesh = _resolve_mesh(mesh)
+    if resolved_mesh is None:
+        raise ValueError("No mesh found")
+    assert not isinstance(resolved_mesh, dict)
 
     def to_sharding(node: typing.Any, spec: typing.Any):
         if spec is None:
@@ -306,7 +352,7 @@ def infer_resource_partitions(
             else:
                 return None
         else:
-            return NamedSharding(mesh, spec)
+            return NamedSharding(resolved_mesh, spec)
 
     return htu.tree_map(to_sharding, tree, pspecs)
 
@@ -407,7 +453,7 @@ class _NamedJitWrapper(eqx.Module):
             if out_axis_resources is not None:
                 # TODO: when AUTO is fixed (or eval_shape can give shardings), use it here
                 out_resources = infer_resource_partitions(
-                    output_shape, out_axis_resources, preserve_existing_shardings=False, use_auto_sharding=False
+                    output_shape, out_axis_resources, preserve_existing_shardings=False
                 )
                 my_pjit_args["out_shardings"] = out_resources
 
@@ -642,7 +688,7 @@ def physical_axis_name(axis: AxisSelector, mapping: Optional[ResourceMapping] =
 def physical_axis_size(axis: AxisSelector, mapping: Optional[ResourceMapping] = None) -> Optional[int]:
     """Get the physical axis size for a logical axis. This is the product of the size of all physical axes
     that this logical axis is mapped to."""
-    mesh = _get_mesh()
+    mesh = _resolve_mesh()
 
     if mesh is None:
         raise ValueError("No mesh found")
@@ -659,10 +705,14 @@ def physical_axis_size(axis: AxisSelector, mapping: Optional[ResourceMapping] =
 
 
 def sharding_for_axis(
-    axis: AxisSelection, mapping: Optional[ResourceMapping] = None, mesh: Optional[Mesh] = None
+    axis: AxisSelection, mapping: Optional[ResourceMapping] = None, mesh: Optional[MeshLike] = None
 ) -> NamedSharding:
     """Get the sharding for a single axis"""
-    return NamedSharding(mesh or _get_mesh(), pspec_for_axis(axis, mapping))
+    resolved_mesh = _resolve_mesh(mesh)
+    if resolved_mesh is None:
+        raise ValueError("No mesh found")
+
+    return NamedSharding(resolved_mesh, pspec_for_axis(axis, mapping))
 
 
 def pspec_for_axis(axis: AxisSelection, mapping: Optional[ResourceMapping] = None) -> PartitionSpec:
@@ -681,28 +731,17 @@ def round_axis_for_partitioning(axis: Axis, mapping: Optional[ResourceMapping] =
         return Axis(axis.name, new_size)
 
 
-def _get_mesh() -> Mesh | AbstractMesh:
-    """Return the current mesh.
-
-    On newer versions of JAX this prefers ``get_abstract_mesh`` which does not
-    capture concrete devices.  If no abstract mesh is currently active we fall
-    back to the concrete mesh used by ``Mesh``'s context manager so existing
-    code continues to work.
-    """
+def _get_mesh() -> Mesh | None:
+    """Deprecated helper that simply proxies to :func:`get_abstract_mesh`."""
 
-    try:  # jax>=0.4.26
-        mesh = get_abstract_mesh()
-        if not getattr(mesh, "empty", False):
-            return mesh
-    except Exception:  # pragma: no cover - older JAX versions
-        pass
-
-    try:
-        from jax.interpreters.pxla import thread_resources
-    except Exception:  # pragma: no cover - jax<0.4
-        from jax.experimental.maps import thread_resources
+    warnings.warn(
+        "`_get_mesh` is deprecated; use `jax's get_abstract_mesh or get_concrete_mesh` instead",
+        DeprecationWarning,
+        stacklevel=2,
+    )
 
-    return thread_resources.env.physical_mesh
+    mesh = _resolve_mesh()
+    return mesh
 
 
 def _is_jit_tracer(x) -> bool: