haliax 1.4.dev419__tar.gz → 1.4.dev438__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev419
+Version: 1.4.dev438
 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,7 +14,7 @@ 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

{haliax-1.4.dev419 → haliax-1.4.dev438}/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",
@@ -60,7 +60,7 @@ haliax = ["src/haliax/*"]
 
 [tool.black]
 line-length = 119
-target-version = ["py310"]
+target-version = ["py311"]
 preview = true
 
 [tool.isort]

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

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

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

@@ -4,7 +4,7 @@
 
 
 import typing as t
-from typing import Optional, Sequence
+from typing import Sequence
 
 import jax
 import jax.numpy as jnp
@@ -139,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,))
@@ -172,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.
 
@@ -208,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.
@@ -226,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.
@@ -238,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.
@@ -260,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)
@@ -587,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
@@ -683,50 +683,50 @@ _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
@@ -735,19 +735,19 @@ def nanmean(
 
 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
@@ -756,11 +756,11 @@ def nanprod(
 
 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
@@ -769,10 +769,10 @@ def nanstd(
 
 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
@@ -781,11 +781,11 @@ def nansum(
 
 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
@@ -795,28 +795,28 @@ def nanvar(
 # "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)
     """

{haliax-1.4.dev419 → haliax-1.4.dev438}/src/haliax/_src/dot.py

@@ -6,7 +6,6 @@
 import functools as ft
 import typing
 import warnings
-from typing import Dict, Optional, Tuple
 
 import jax
 
@@ -29,11 +28,11 @@ from haliax.types import DTypeLike, PrecisionLike
 # deprecated overload
 @typing.overload
 def dot(
-    axis: Optional[AxisSelection],
+    axis: AxisSelection | None,
     *arrays: NamedArray,
     precision: PrecisionLike = None,
-    preferred_element_type: Optional[DTypeLike] = None,
-    out_axes: Optional[PartialAxisSpec] = ...,
+    preferred_element_type: DTypeLike | None = None,
+    out_axes: PartialAxisSpec | None = ...,
     dot_general=jax.lax.dot_general,
 ) -> NamedArray: ...
 
@@ -41,10 +40,10 @@ def dot(
 @typing.overload
 def dot(
     *arrays: NamedArray,
-    axis: Optional[AxisSelection],
+    axis: AxisSelection | None,
     precision: PrecisionLike = None,
-    preferred_element_type: Optional[DTypeLike] = None,
-    out_axes: Optional[PartialAxisSpec] = ...,
+    preferred_element_type: DTypeLike | None = None,
+    out_axes: PartialAxisSpec | None = ...,
     dot_general=jax.lax.dot_general,
 ) -> NamedArray: ...
 
@@ -52,8 +51,8 @@ def dot(
 def dot(
     *arrays,
     precision: PrecisionLike = None,
-    preferred_element_type: Optional[DTypeLike] = None,
-    out_axes: Optional[PartialAxisSpec] = None,
+    preferred_element_type: DTypeLike | None = None,
+    out_axes: PartialAxisSpec | None = None,
     dot_general=jax.lax.dot_general,
     **kwargs,
 ) -> NamedArray:
@@ -82,7 +81,7 @@ def dot(
             which in turn passes it to jax.lax.dot_general.
         preferred_element_type (DTypeLike, optional): The preferred element type of the result. Defaults to None.
             This argument is passed to `jax.numpy.einsum`.
-        out_axes (Optional[PartialAxisSpec], optional): a potentially partial specification of the output axes.
+        out_axes (PartialAxisSpec | None, optional): a potentially partial specification of the output axes.
             If provided, the output will be transposed to match the provided axes. Defaults to None.
 
 
@@ -107,8 +106,8 @@ def dot(
     # to call dot_general we need two things:
     # list of contractions and list of arrays
 
-    all_axes: Tuple[Axis, ...] = ft.reduce(union_axes, (a.axes for a in arrays), ())  # type: ignore
-    output_axes: Tuple[Axis, ...]
+    all_axes: tuple[Axis, ...] = ft.reduce(union_axes, (a.axes for a in arrays), ())  # type: ignore
+    output_axes: tuple[Axis, ...]
     if axis is None:
         # we want to contract over all the axes
         output_axes = ()
@@ -121,7 +120,7 @@ def dot(
     array_specs = []
 
     next_index = 0
-    axis_mappings: Dict[str, int] = {}
+    axis_mappings: dict[str, int] = {}
 
     for a in arrays:
         spec = ""

{haliax-1.4.dev419 → haliax-1.4.dev438}/src/haliax/_src/einsum.py

@@ -5,7 +5,6 @@
 
 import functools
 from types import EllipsisType
-from typing import Optional, Tuple
 
 import jax.lax
 
@@ -24,7 +23,7 @@ def einsum(
     equation: str,
     *arrays: NamedArray,
     precision: PrecisionLike = None,
-    preferred_element_type: Optional[DTypeLike] = None,
+    preferred_element_type: DTypeLike | None = None,
     _dot_general: DotGeneralOp = jax.lax.dot_general,
     **axis_aliases: AxisSelector,
 ) -> NamedArray:
@@ -306,7 +305,7 @@ def _all_input_axes(arrays):
     return ensure_tuple(functools.reduce(union_axes, (a.axes for a in arrays), ()))  # type: ignore
 
 
-def _captures_to_axis_names(equation, lhs, aliases) -> Tuple[list[str | EllipsisType], bool, set[str]]:
+def _captures_to_axis_names(equation, lhs, aliases) -> tuple[list[str | EllipsisType], bool, set[str]]:
     covered_aliases = set()
     candidate_axes: list[str | EllipsisType] = []
     has_ellipsis = False

{haliax-1.4.dev419 → haliax-1.4.dev438}/src/haliax/_src/parsing.py

@@ -5,16 +5,16 @@
 
 import dataclasses
 from types import EllipsisType
-from typing import Mapping, NoReturn, Optional, Sequence
+from typing import Mapping, NoReturn, Sequence
 
 from haliax.axis import Axis, AxisSelector
 
 
 @dataclasses.dataclass(frozen=True)
 class _AxisCapture:
-    binding: Optional[str] = None
+    binding: str | None = None
     axes: tuple[str, ...] = ()
-    char_range: Optional[tuple[int, int]] = None
+    char_range: tuple[int, int] | None = None
 
     def __post_init__(self):
         if len(self.axes) == 0:
@@ -27,7 +27,7 @@ class Expression:
     is_ordered: bool
 
 
-def raise_parse_error(message: str, expression: str, pos: Optional[int | tuple[int, int]]) -> NoReturn:
+def raise_parse_error(message: str, expression: str, pos: int | tuple[int, int] | None) -> NoReturn:
     """Raise a ValueError with a message and the position in the expression."""
     fmt = f"Error while parsing:\n    {expression}"
     if pos is not None:
@@ -234,7 +234,7 @@ class AliasTable:
         else:
             self.bindings = {**bindings}
 
-    def dealias_binding(self, binding: str) -> Optional[AxisSelector]:
+    def dealias_binding(self, binding: str) -> AxisSelector | None:
         return self.bindings.get(binding, None)
 
     def bind_alias(self, alias: str, axis: Axis, expr, char_range):

{haliax-1.4.dev419 → haliax-1.4.dev438}/src/haliax/_src/rearrange.py

@@ -7,7 +7,7 @@
 import dataclasses
 import typing
 from types import EllipsisType
-from typing import Mapping, Optional, Sequence
+from typing import Mapping, Sequence
 
 import jax.lax
 import jax.numpy as jnp
@@ -157,7 +157,7 @@ def einops_rearrange(array: NamedArray, expression: str, **bindings: AxisSelecto
 @dataclasses.dataclass(frozen=True)
 class _Plan:
     intermediate_axes: tuple[Axis, ...]
-    transpose: Optional[tuple[int, ...]]
+    transpose: tuple[int, ...] | None
     needs_final_reshape: bool
 
     final_axes: tuple[Axis, ...]
@@ -170,7 +170,7 @@ def _plan_rearrange(
     grouped_new_shapes = _determine_initial_reshape(original_str, lhs, array, aliases)
     intermediate_axes = tuple(ax for split_axes in grouped_new_shapes for ax in split_axes)
 
-    transpose: Optional[tuple[int, ...]]
+    transpose: tuple[int, ...] | None
     transpose, final_axes = _determine_final_transpose_and_reshape(original_str, rhs, aliases, intermediate_axes)
 
     transposed_intermediate_axes = tuple(intermediate_axes[i] for i in transpose)
@@ -289,7 +289,7 @@ def _determine_initial_reshape(
     # the lhs all need to be bound to axes in the array, or synthesized as parts of axes.
     # In the lhs, bindings look like either a name, or a name and a list of (new) axes.
     # bindings can either be done by name, or by position, depending on if lhs.is_ordered
-    new_shapes: list[Optional[list[Axis]]] = [None] * len(array.axes)
+    new_shapes: list[list[Axis] | None] = [None] * len(array.axes)
     used_new_names: set[str] = set()  # names can only be used once on a side
 
     # one subtle difference between the lhs and the rhs is the handling of binding in expressions like (a: b c)
@@ -301,7 +301,7 @@ def _determine_initial_reshape(
         # if we start with an ellipsis, we bind from the right
         # if we end with an ellipsis, we bind from the left
         ellipsis_pos = None
-        axis_index_for_capture: list[Optional[int]] = [None] * len(lhs.captures)
+        axis_index_for_capture: list[int | None] = [None] * len(lhs.captures)
         covered_axes = set()
         # bind from the left
         axis_pos = 0
@@ -414,8 +414,8 @@ def _solve_split_axes(axis, capture, aliases, used_new_names, expression):
     """
     Given an axis and a capture of the form (a: b c) or (b c) on the lhs, solve for the new axes.
     """
-    new_axes: list[Optional[Axis]] = []
-    unsolved_axis_index: Optional[int] = None
+    new_axes: list[Axis | None] = []
+    unsolved_axis_index: int | None = None
 
     # easy case: 1 axis in capture
     if len(capture.axes) == 1: