haliax 1.4.dev413__tar.gz → 1.4.dev420__tar.gz

{haliax-1.4.dev413 → haliax-1.4.dev420}/.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.dev420}/.github/workflows/publish_dev.yaml

@@ -20,7 +20,7 @@ jobs:
       - name: Set up Python
         uses: actions/setup-python@v2
         with:
-          python-version: '3.x'
+          python-version: '3.11'
 
       - name: Calculate Version and Build Number
         run: |

{haliax-1.4.dev413 → haliax-1.4.dev420}/.github/workflows/run_pre_commit.yaml

@@ -8,7 +8,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.10.11"]
+        python-version: ["3.11"]
 
     steps:
       - uses: actions/checkout@v3

{haliax-1.4.dev413 → haliax-1.4.dev420}/.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.11"]
     steps:
       - name: Checkout repository
         uses: actions/checkout@v3

{haliax-1.4.dev413 → haliax-1.4.dev420}/.github/workflows/run_tests.yaml

@@ -9,10 +9,10 @@ jobs:
 
     steps:
       - uses: actions/checkout@v3
-      - name: Set up Python 3.10.11
+      - name: Set up Python 3.11
         uses: actions/setup-python@v4
         with:
-          python-version: 3.10.11
+          python-version: 3.11
       - name: Install dependencies
         run: |
           python -m pip install uv

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev413
+Version: 1.4.dev420
 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/
@@ -14,9 +14,10 @@ Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: MacOS :: MacOS X
 Classifier: Operating System :: POSIX :: Linux
 Classifier: Programming Language :: Python :: 3
-Requires-Python: >=3.10
+Requires-Python: >=3.11
 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.dev420}/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.dev420}/pyproject.toml

@@ -11,7 +11,7 @@ authors = [
 ]
 description = "Named Tensors for Legible Deep Learning in JAX"
 readme = "README.md"
-requires-python = ">=3.10"
+requires-python = ">=3.11"
 classifiers = [
     "Programming Language :: Python :: 3",
     "License :: OSI Approved :: Apache Software License",
@@ -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",
@@ -61,7 +60,7 @@ haliax = ["src/haliax/*"]
 
 [tool.black]
 line-length = 119
-target-version = ["py310"]
+target-version = ["py311"]
 preview = true
 
 [tool.isort]

{haliax-1.4.dev413 → haliax-1.4.dev420}/src/haliax/__about__.py

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

{haliax-1.4.dev413 → haliax-1.4.dev420}/src/haliax/__init__.py

@@ -4,15 +4,12 @@
 
 
 import typing as t
-from typing import Optional, Sequence
+from typing import 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
@@ -126,21 +139,21 @@ A = t.TypeVar("A", Scalar, NamedArray, jnp.ndarray)
 
 
 # creation routines
-def zeros(shape: AxisSpec, dtype: Optional[DTypeLike] = None) -> NamedArray:
+def zeros(shape: AxisSpec, dtype: DTypeLike | None = None) -> NamedArray:
     """Creates a NamedArray with all elements set to 0"""
     if dtype is None:
         dtype = jnp.float32
     return full(shape, 0, dtype)
 
 
-def ones(shape: AxisSpec, dtype: Optional[DTypeLike] = None) -> NamedArray:
+def ones(shape: AxisSpec, dtype: DTypeLike | None = None) -> NamedArray:
     """Creates a NamedArray with all elements set to 1"""
     if dtype is None:
         dtype = jnp.float32
     return full(shape, 1, dtype)
 
 
-def full(shape: AxisSpec, fill_value: T, dtype: Optional[DTypeLike] = None) -> NamedArray:
+def full(shape: AxisSpec, fill_value: T, dtype: DTypeLike | None = None) -> NamedArray:
     """Creates a NamedArray with all elements set to `fill_value`"""
     if isinstance(shape, Axis):
         return NamedArray(jnp.full(shape=shape.size, fill_value=fill_value, dtype=dtype), (shape,))
@@ -159,12 +172,12 @@ def ones_like(a: NamedArray, dtype=None) -> NamedArray:
     return NamedArray(jnp.ones_like(a.array, dtype=dtype), a.axes)
 
 
-def full_like(a: NamedArray, fill_value: T, dtype: Optional[DTypeLike] = None) -> NamedArray:
+def full_like(a: NamedArray, fill_value: T, dtype: DTypeLike | None = None) -> NamedArray:
     """Creates a NamedArray with all elements set to `fill_value`"""
     return NamedArray(jnp.full_like(a.array, fill_value, dtype=dtype), a.axes)
 
 
-def arange(axis: AxisSpec, *, start=0, step=1, dtype: Optional[DTypeLike] = None) -> NamedArray:
+def arange(axis: AxisSpec, *, start=0, step=1, dtype: DTypeLike | None = None) -> NamedArray:
     """
     Version of jnp.arange that returns a NamedArray.
 
@@ -195,7 +208,7 @@ def arange(axis: AxisSpec, *, start=0, step=1, dtype: Optional[DTypeLike] = None
 
 # TODO: add overrides for arraylike start/stop to linspace, logspace, geomspace
 def linspace(
-    axis: AxisSelector, *, start: float, stop: float, endpoint: bool = True, dtype: Optional[DTypeLike] = None
+    axis: AxisSelector, *, start: float, stop: float, endpoint: bool = True, dtype: DTypeLike | None = None
 ) -> NamedArray:
     """
     Version of jnp.linspace that returns a NamedArray.
@@ -213,7 +226,7 @@ def logspace(
     stop: float,
     endpoint: bool = True,
     base: float = 10.0,
-    dtype: Optional[DTypeLike] = None,
+    dtype: DTypeLike | None = None,
 ) -> NamedArray:
     """
     Version of jnp.logspace that returns a NamedArray.
@@ -225,7 +238,7 @@ def logspace(
 
 
 def geomspace(
-    axis: AxisSelector, *, start: float, stop: float, endpoint: bool = True, dtype: Optional[DTypeLike] = None
+    axis: AxisSelector, *, start: float, stop: float, endpoint: bool = True, dtype: DTypeLike | None = None
 ) -> NamedArray:
     """
     Version of jnp.geomspace that returns a NamedArray.
@@ -247,7 +260,7 @@ def stack(axis: AxisSelector, arrays: Sequence[NamedArray]) -> NamedArray:
 
 
 def repeat(
-    a: NamedArray, repeats: int | jnp.ndarray, axis: AxisSelector, total_repeat_length: Optional[int] = None
+    a: NamedArray, repeats: int | jnp.ndarray, axis: AxisSelector, total_repeat_length: int | None = None
 ) -> NamedArray:
     """Version of [jax.numpy.repeat][] that returns a NamedArray"""
     index = a.axis_indices(axis)
@@ -574,91 +587,91 @@ def trunc(a: A) -> A:
 
 
 # Reduction functions
-def all(array: NamedArray, axis: Optional[AxisSelection] = None, *, where: Optional[NamedArray] = None) -> NamedArray:
+def all(array: NamedArray, axis: AxisSelection | None = None, *, where: NamedArray | None = None) -> NamedArray:
     """
     Named version of [jax.numpy.all](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.all.html#jax.numpy.all).
     """
     return wrap_reduction_call(jnp.all, array, axis, where, single_axis_only=False, supports_where=True)
 
 
-def amax(array: NamedArray, axis: Optional[AxisSelection] = None, *, where: Optional[NamedArray] = None) -> NamedArray:
+def amax(array: NamedArray, axis: AxisSelection | None = None, *, where: NamedArray | None = None) -> NamedArray:
     """
     Aliax for max. See max for details.
     """
     return wrap_reduction_call(jnp.amax, array, axis, where, single_axis_only=False, supports_where=True)
 
 
-def amin(array: NamedArray, axis: Optional[AxisSelection] = None, *, where: Optional[NamedArray] = None) -> NamedArray:
+def amin(array: NamedArray, axis: AxisSelection | None = None, *, where: NamedArray | None = None) -> NamedArray:
     """
     Aliax for min. See min for details.
     """
     return wrap_reduction_call(jnp.amin, array, axis, where, single_axis_only=False, supports_where=True)
 
 
-def any(array: NamedArray, axis: Optional[AxisSelection] = None, *, where: Optional[NamedArray] = None) -> NamedArray:
+def any(array: NamedArray, axis: AxisSelection | None = None, *, where: NamedArray | None = None) -> NamedArray:
     """True if any elements along a given axis or axes are True. If axis is None, any elements are True."""
     return wrap_reduction_call(jnp.any, array, axis, where, single_axis_only=False, supports_where=True)
 
 
-def argmax(array: NamedArray, axis: Optional[AxisSelector]) -> NamedArray:
+def argmax(array: NamedArray, axis: AxisSelector | None) -> NamedArray:
     return wrap_reduction_call(jnp.argmax, array, axis, None, single_axis_only=True, supports_where=False)
 
 
-def argmin(array: NamedArray, axis: Optional[AxisSelector]) -> NamedArray:
+def argmin(array: NamedArray, axis: AxisSelector | None) -> NamedArray:
     return wrap_reduction_call(jnp.argmin, array, axis, None, single_axis_only=True, supports_where=False)
 
 
-def max(array: NamedArray, axis: Optional[AxisSelection] = None, *, where: Optional[NamedArray] = None) -> NamedArray:
+def max(array: NamedArray, axis: AxisSelection | None = None, *, where: NamedArray | None = None) -> NamedArray:
     return wrap_reduction_call(jnp.max, array, axis, where, single_axis_only=False, supports_where=True)
 
 
 def mean(
     array: NamedArray,
-    axis: Optional[AxisSelection] = None,
+    axis: AxisSelection | None = None,
     *,
-    where: Optional[NamedArray] = None,
-    dtype: Optional[DTypeLike] = None,
+    where: NamedArray | None = None,
+    dtype: DTypeLike | None = None,
 ) -> NamedArray:
     return wrap_reduction_call(jnp.mean, array, axis, where, single_axis_only=False, supports_where=True, dtype=dtype)
 
 
-def min(array: NamedArray, axis: Optional[AxisSelection] = None, *, where: Optional[NamedArray] = None) -> NamedArray:
+def min(array: NamedArray, axis: AxisSelection | None = None, *, where: NamedArray | None = None) -> NamedArray:
     return wrap_reduction_call(jnp.min, array, axis, where, single_axis_only=False, supports_where=True)
 
 
 def prod(
     array: NamedArray,
-    axis: Optional[AxisSelection] = None,
+    axis: AxisSelection | None = None,
     *,
-    where: Optional[NamedArray] = None,
-    dtype: Optional[DTypeLike] = None,
+    where: NamedArray | None = None,
+    dtype: DTypeLike | None = None,
 ) -> NamedArray:
     return wrap_reduction_call(jnp.prod, array, axis, where, single_axis_only=False, supports_where=True, dtype=dtype)
 
 
 def std(
     array: NamedArray,
-    axis: Optional[AxisSelection] = None,
+    axis: AxisSelection | None = None,
     *,
-    where: Optional[NamedArray] = None,
+    where: NamedArray | None = None,
     ddof: int = 0,
-    dtype: Optional[DTypeLike] = None,
+    dtype: DTypeLike | None = None,
 ) -> NamedArray:
     return wrap_reduction_call(
         jnp.std, array, axis, where, single_axis_only=False, supports_where=True, dtype=dtype, ddof=ddof
     )
 
 
-def ptp(array: NamedArray, axis: Optional[AxisSelection] = None, *, where: Optional[NamedArray] = None) -> NamedArray:
+def ptp(array: NamedArray, axis: AxisSelection | None = None, *, where: NamedArray | None = None) -> NamedArray:
     return wrap_reduction_call(jnp.ptp, array, axis, where, single_axis_only=False, supports_where=True)
 
 
 def product(
     array: NamedArray,
-    axis: Optional[AxisSelection] = None,
+    axis: AxisSelection | None = None,
     *,
-    where: Optional[NamedArray] = None,
-    dtype: Optional[DTypeLike] = None,
+    where: NamedArray | None = None,
+    dtype: DTypeLike | None = None,
 ) -> NamedArray:
     return wrap_reduction_call(
         jnp.product, array, axis, where, single_axis_only=False, supports_where=True, dtype=dtype
@@ -670,130 +683,140 @@ _sum = sum
 
 def sum(
     array: NamedArray,
-    axis: Optional[AxisSelection] = None,
+    axis: AxisSelection | None = None,
     *,
-    where: Optional[NamedArray] = None,
-    dtype: Optional[DTypeLike] = None,
+    where: NamedArray | None = None,
+    dtype: DTypeLike | None = None,
 ) -> NamedArray:
     return wrap_reduction_call(jnp.sum, array, axis, where, single_axis_only=False, supports_where=True, dtype=dtype)
 
 
 def var(
     array: NamedArray,
-    axis: Optional[AxisSelection] = None,
+    axis: AxisSelection | None = None,
     *,
-    where: Optional[NamedArray] = None,
+    where: NamedArray | None = None,
     ddof: int = 0,
-    dtype: Optional[DTypeLike] = None,
+    dtype: DTypeLike | None = None,
 ) -> NamedArray:
     return wrap_reduction_call(
         jnp.var, array, axis, where, single_axis_only=False, supports_where=True, dtype=dtype, ddof=ddof
     )
 
 
-def nanargmax(array: NamedArray, axis: Optional[AxisSelector] = None) -> NamedArray:
+def nanargmax(array: NamedArray, axis: AxisSelector | None = None) -> NamedArray:
     return wrap_reduction_call(jnp.nanargmax, array, axis, None, single_axis_only=True, supports_where=False)
 
 
-def nanargmin(array: NamedArray, axis: Optional[AxisSelector] = None) -> NamedArray:
+def nanargmin(array: NamedArray, axis: AxisSelector | None = None) -> NamedArray:
     return wrap_reduction_call(jnp.nanargmin, array, axis, None, single_axis_only=True, supports_where=False)
 
 
 def nanmax(
     array: NamedArray,
-    axis: Optional[AxisSelection] = None,
+    axis: AxisSelection | None = None,
     *,
-    where: Optional[NamedArray] = None,
+    where: NamedArray | None = None,
 ) -> NamedArray:
     return wrap_reduction_call(jnp.nanmax, array, axis, where, single_axis_only=False, supports_where=True)
 
 
 def nanmean(
     array: NamedArray,
-    axis: Optional[AxisSelection] = None,
+    axis: AxisSelection | None = None,
     *,
-    where: Optional[NamedArray] = None,
-    dtype: Optional[DTypeLike] = None,
+    where: NamedArray | None = None,
+    dtype: DTypeLike | None = 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(
     array: NamedArray,
-    axis: Optional[AxisSelection] = None,
+    axis: AxisSelection | None = None,
     *,
-    where: Optional[NamedArray] = None,
+    where: NamedArray | None = None,
 ) -> NamedArray:
     return wrap_reduction_call(jnp.nanmin, array, axis, where, single_axis_only=False, supports_where=True)
 
 
 def nanprod(
     array: NamedArray,
-    axis: Optional[AxisSelection] = None,
+    axis: AxisSelection | None = None,
     *,
-    where: Optional[NamedArray] = None,
-    dtype: Optional[DTypeLike] = None,
+    where: NamedArray | None = None,
+    dtype: DTypeLike | None = 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(
     array: NamedArray,
-    axis: Optional[AxisSelection] = None,
+    axis: AxisSelection | None = None,
     *,
-    where: Optional[NamedArray] = None,
+    where: NamedArray | None = None,
     ddof: int = 0,
-    dtype: Optional[DTypeLike] = None,
+    dtype: DTypeLike | None = 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(
     array: NamedArray,
-    axis: Optional[AxisSelection] = None,
+    axis: AxisSelection | None = None,
     *,
-    where: Optional[NamedArray] = None,
-    dtype: Optional[DTypeLike] = None,
+    where: NamedArray | None = None,
+    dtype: DTypeLike | None = 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(
     array: NamedArray,
-    axis: Optional[AxisSelection] = None,
+    axis: AxisSelection | None = None,
     *,
-    where: Optional[NamedArray] = None,
+    where: NamedArray | None = None,
     ddof: int = 0,
-    dtype: Optional[DTypeLike] = None,
+    dtype: DTypeLike | None = 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
 
 
-def cumsum(a: NamedArray, axis: AxisSelector, *, dtype: Optional[DTypeLike] = None) -> NamedArray:
+def cumsum(a: NamedArray, axis: AxisSelector, *, dtype: DTypeLike | None = None) -> NamedArray:
     """
     Named version of [jax.numpy.cumsum](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.cumsum.html)
     """
     return wrap_axiswise_call(jnp.cumsum, a, axis, dtype=dtype, single_axis_only=True)
 
 
-def cumprod(a: NamedArray, axis: AxisSelector, dtype: Optional[DTypeLike] = None) -> NamedArray:
+def cumprod(a: NamedArray, axis: AxisSelector, dtype: DTypeLike | None = None) -> NamedArray:
     """
     Named version of [jax.numpy.cumprod](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.cumprod.html)
     """
     return wrap_axiswise_call(jnp.cumprod, a, axis, dtype=dtype, single_axis_only=True)
 
 
-def nancumsum(a: NamedArray, axis: AxisSelector, *, dtype: Optional[DTypeLike] = None) -> NamedArray:
+def nancumsum(a: NamedArray, axis: AxisSelector, *, dtype: DTypeLike | None = None) -> NamedArray:
     """
     Named version of [jax.numpy.nancumsum](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.nancumsum.html)
     """
     return wrap_axiswise_call(jnp.nancumsum, a, axis, dtype=dtype, single_axis_only=True)
 
 
-def nancumprod(a: NamedArray, axis: AxisSelector, dtype: Optional[DTypeLike] = None) -> NamedArray:
+def nancumprod(a: NamedArray, axis: AxisSelector, dtype: DTypeLike | None = None) -> NamedArray:
     """
     Named version of [jax.numpy.nancumprod](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.nancumprod.html)
     """
@@ -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",
 ]