haliax 1.4.dev409__tar.gz → 1.4.dev411__tar.gz

{haliax-1.4.dev409 → haliax-1.4.dev411}/.agents/projects/api_parity.md

@@ -107,7 +107,7 @@ APIs that don't translate well to named tensors are intentionally omitted here.
 - [x] `nanstd`
 - [x] `nansum`
 - [x] `nanvar`
-- [ ] `nonzero`
+- [x] `nonzero`
 - [ ] `ogrid`
 - [x] `packbits`
 - [ ] `partition`

haliax-1.4.dev411/.pre-commit-config.yaml

@@ -0,0 +1,43 @@
+# See https://pre-commit.com
+exclude: ".git|.venv|tests/snapshots/.*/.*"
+
+repos:
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.11.10
+    hooks:
+      - id: ruff
+        args: [ --fix, --exit-non-zero-on-fix ]
+
+  - repo: https://github.com/Lucas-C/pre-commit-hooks
+    rev: v1.5.5
+    hooks:
+      - id: insert-license
+        files: \.py$
+        args:
+          - --license-filepath
+          - etc/license_header.txt
+          - --use-current-year
+
+  - repo: https://github.com/psf/black
+    rev: 25.1.0
+    hooks:
+      - id: black
+
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v5.0.0
+    hooks:
+      - id: check-added-large-files
+      - id: check-ast
+      - id: check-case-conflict
+      - id: check-merge-conflict
+      - id: check-toml
+      - id: check-yaml
+        args: [ --unsafe ]
+      - id: end-of-file-fixer
+      - id: trailing-whitespace
+
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v1.16.1
+    hooks:
+      - id: mypy
+        args: [--ignore-missing-imports, --check-untyped-defs]

haliax-1.4.dev411/AUTHORS.md

@@ -0,0 +1,5 @@
+# The Levanter Authors
+
+The Levanter Authors currently include:
+
+- The Board of Trustees of the Leland Stanford Junior University

haliax-1.4.dev411/CONTRIBUTORS.md

@@ -0,0 +1,15 @@
+# Contributors
+
+The following individuals have contributed to Haliax:
+
+- David Hall <dlwh@stanford.edu>
+- David Hall <dlwh@cs.stanford.edu>
+- Jason Wang <blahblahj.wsy@gmail.com>
+- Ivan Zhou <ivan.zhouyq@gmail.com>
+- rohan-mehta-1024 <69774557+rohan-mehta-1024@users.noreply.github.com>
+- Gary Miguel <garymm@garymm.org>
+- Jennifer Zhou <jennifer@jezh.me>
+- Joseph Camacho <camacho.joseph@gmail.com>
+- Omead Pooladzandi <opooladz@ucla.edu>
+- Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com>
+- Russell Power <russell.power@gmail.com>

{haliax-1.4.dev409 → haliax-1.4.dev411}/PKG-INFO

@@ -1,11 +1,12 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev409
+Version: 1.4.dev411
 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/
 Project-URL: Documentation, https://haliax.readthedocs.io/en/latest/
 Author-email: David Hall <dlwh@cs.stanford.edu>
+License-File: AUTHORS.md
 License-File: LICENSE
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Science/Research

{haliax-1.4.dev409 → haliax-1.4.dev411}/docs/api.md

@@ -283,6 +283,7 @@ These are all more or less directly from JAX's NumPy API.
 ::: haliax.pad
 ::: haliax.searchsorted
 ::: haliax.top_k
+::: haliax.nonzero
 ::: haliax.trace
 ::: haliax.tril
 ::: haliax.triu

haliax-1.4.dev411/etc/license_header.txt

@@ -0,0 +1,3 @@
+Copyright 2025 The Levanter Authors
+
+SPDX-License-Identifier: Apache-2.0

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

@@ -0,0 +1,6 @@
+# Copyright 2025 The Levanter Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+
+__version__ = "1.4.dev411"

{haliax-1.4.dev409 → haliax-1.4.dev411}/src/haliax/__init__.py

@@ -1,10 +1,14 @@
+# Copyright 2025 The Levanter Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+
 import typing as t
 from typing import Optional, Sequence
 
 import jax
 import jax.numpy as jnp
 
-
 try:
     from jax.typing import DTypeLike
 except ImportError:
@@ -44,7 +48,9 @@ from .axis import (
 )
 from .core import (
     NamedArray,
-    NamedArrayAxes, NamedArrayAxesSpec, NamedOrNumeric,
+    NamedArrayAxes,
+    NamedArrayAxesSpec,
+    NamedOrNumeric,
     are_shape_checks_enabled,
     broadcast_arrays,
     broadcast_axis,
@@ -78,6 +84,7 @@ from .ops import (
     trace,
     tril,
     triu,
+    nonzero,
     unique,
     unique_values,
     unique_counts,
@@ -102,7 +109,6 @@ from .wrap import (
     wrap_reduction_call,
 )
 
-
 T = t.TypeVar("T")
 A = t.TypeVar("A", Scalar, NamedArray, jnp.ndarray)
 
@@ -801,6 +807,7 @@ def argsort(a: NamedArray, axis: AxisSelector) -> NamedArray:
 
 # elemwise binary ops
 
+
 # Note that all the heavy lifting is done by the `wrap_elemwise_binary` decorator
 @wrap_elemwise_binary
 def add(x1: NamedOrNumeric, x2: NamedOrNumeric, /) -> NamedOrNumeric:
@@ -1193,6 +1200,7 @@ __all__ = [
     "map",
     "vmap",
     "trace",
+    "nonzero",
     "where",
     "unique",
     "unique_values",

haliax-1.4.dev411/src/haliax/_src/__init__.py

@@ -0,0 +1,3 @@
+# Copyright 2025 The Levanter Authors
+#
+# SPDX-License-Identifier: Apache-2.0

{haliax-1.4.dev409 → haliax-1.4.dev411}/src/haliax/_src/compile_utils.py

@@ -1,3 +1,8 @@
+# Copyright 2025 The Levanter Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+
 # This whole file is copied from Equinox.
 # (c) 2023, Google LLC. and/or Patrick Kidger. Apache 2.0 licensed.
 # Patrick doesn't like that I depend on Equinox internals, so I copied this stuff

{haliax-1.4.dev409 → haliax-1.4.dev411}/src/haliax/_src/dot.py

@@ -1,3 +1,8 @@
+# Copyright 2025 The Levanter Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+
 import functools as ft
 import typing
 import warnings
@@ -30,8 +35,7 @@ def dot(
     preferred_element_type: Optional[DTypeLike] = None,
     out_axes: Optional[PartialAxisSpec] = ...,
     dot_general=jax.lax.dot_general,
-) -> NamedArray:
-    ...
+) -> NamedArray: ...
 
 
 @typing.overload
@@ -42,8 +46,7 @@ def dot(
     preferred_element_type: Optional[DTypeLike] = None,
     out_axes: Optional[PartialAxisSpec] = ...,
     dot_general=jax.lax.dot_general,
-) -> NamedArray:
-    ...
+) -> NamedArray: ...
 
 
 def dot(

{haliax-1.4.dev409 → haliax-1.4.dev411}/src/haliax/_src/einsum.py

@@ -1,3 +1,8 @@
+# Copyright 2025 The Levanter Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+
 import functools
 from types import EllipsisType
 from typing import Optional, Tuple

{haliax-1.4.dev409 → haliax-1.4.dev411}/src/haliax/_src/fp8.py

@@ -1,9 +1,13 @@
+# Copyright 2025 The Levanter Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+
 from functools import partial
 
 from jax import custom_jvp, custom_vjp, lax
 from jax import numpy as jnp
 
-
 # All of this is copy paste from flax/linen/fp8_ops.py
 # (Until we get to the module)
 

{haliax-1.4.dev409 → haliax-1.4.dev411}/src/haliax/_src/parsing.py

@@ -1,3 +1,8 @@
+# Copyright 2025 The Levanter Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+
 import dataclasses
 from types import EllipsisType
 from typing import Mapping, NoReturn, Optional, Sequence

{haliax-1.4.dev409 → haliax-1.4.dev411}/src/haliax/_src/rearrange.py

@@ -1,3 +1,8 @@
+# Copyright 2025 The Levanter Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+
 # Support for einops-style rearrangement strings, but supporting named axes and unordered matching
 import dataclasses
 import typing

{haliax-1.4.dev409 → haliax-1.4.dev411}/src/haliax/_src/scan.py

@@ -1,3 +1,8 @@
+# Copyright 2025 The Levanter Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+
 import dataclasses
 import functools as ft
 import inspect
@@ -16,7 +21,6 @@ from haliax.core import NamedArray
 from haliax.jax_utils import is_jax_array_like, multilevel_scan, tree_checkpoint_name
 from haliax.util import is_jax_or_hax_array_like, is_named_array
 
-
 BoolAxisSpec = Union[bool, Callable[[Any], bool]]
 Carry = TypeVar("Carry")
 X = TypeVar("X", contravariant=True)
@@ -31,8 +35,7 @@ def is_named_or_shaped_array_like(x):
 class ScanFn(Protocol[Carry, Args, Y]):
     """ """
 
-    def __call__(self, carry: Carry, *args: Args.args, **kwargs: Args.kwargs) -> tuple[Carry, Y]:
-        ...
+    def __call__(self, carry: Carry, *args: Args.args, **kwargs: Args.kwargs) -> tuple[Carry, Y]: ...
 
 
 @dataclasses.dataclass(frozen=True)
@@ -262,8 +265,7 @@ def scan(
     reverse: bool = False,
     unroll: int = 1,
     is_scanned: BoolAxisSpec = is_named_or_shaped_array_like,
-) -> Callable[[Carry, PyTree[X]], tuple[Carry, PyTree[Y]]]:
-    ...
+) -> Callable[[Carry, PyTree[X]], tuple[Carry, PyTree[Y]]]: ...
 
 
 @overload
@@ -275,8 +277,7 @@ def scan(
     reverse: bool = False,
     unroll: int = 1,
     is_scanned: BoolAxisSpec = is_named_or_shaped_array_like,
-) -> Callable:
-    ...
+) -> Callable: ...
 
 
 def scan(
@@ -444,8 +445,7 @@ def fold(
     reverse: bool = False,
     unroll: int = 1,
     is_scanned: BoolAxisSpec = is_named_or_shaped_array_like,
-) -> Callable[[Carry, PyTree[X]], Carry]:
-    ...
+) -> Callable[[Carry, PyTree[X]], Carry]: ...
 
 
 @overload
@@ -457,8 +457,7 @@ def fold(
     reverse: bool = False,
     unroll: int = 1,
     is_scanned: BoolAxisSpec = is_named_or_shaped_array_like,
-) -> Callable:
-    ...
+) -> Callable: ...
 
 
 def fold(
@@ -549,9 +548,7 @@ def _zero_if_array_else_none(x: Any) -> ResolvedUnnamedAxisSpec:
     return 0 if is_jax_array_like(x) else None
 
 
-def _format_tree_path(
-    path: tuple[jtu.KeyEntry, ...], arg_pos_names: dict[int, str] | None = None
-) -> str:
+def _format_tree_path(path: tuple[jtu.KeyEntry, ...], arg_pos_names: dict[int, str] | None = None) -> str:
     parts: list[str] = []
     i = 0
     if len(path) >= 2 and isinstance(path[0], jtu.SequenceKey):

{haliax-1.4.dev409 → haliax-1.4.dev411}/src/haliax/_src/state_dict.py

@@ -1,3 +1,8 @@
+# Copyright 2025 The Levanter Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+
 # Module to support torch-style "state dict" serialization via safetensors
 import dataclasses
 import typing
@@ -19,7 +24,6 @@ from haliax.core import NamedArray, named
 from haliax.jax_utils import is_jax_array_like, is_scalarish
 from haliax.tree_util import scan_aware_tree_map
 
-
 try:
     import safetensors
 except ImportError:
@@ -92,6 +96,7 @@ def _flatten_to_unflatten(t, state_dict, prefix):
     """
     Flatten the torch compatible state_dict before loading into t, and then recover the unflattened layers.
     """
+
     # typically, `t` is a bunch of ShapeDtypeStructs, which can't be transposed etc. so we instead have to zeros()
     # into real arrays (that aren't actually real b/c this is inside a jit)
     def _dt_struct_to_array(struct):
@@ -107,18 +112,15 @@ def _flatten_to_unflatten(t, state_dict, prefix):
 
 
 @typing.overload
-def with_prefix(prefix: str | None, leaf: str) -> str:
-    ...
+def with_prefix(prefix: str | None, leaf: str) -> str: ...
 
 
 @typing.overload
-def with_prefix(prefix: str, leaf: None) -> str:
-    ...
+def with_prefix(prefix: str, leaf: None) -> str: ...
 
 
 @typing.overload
-def with_prefix(prefix: Optional[str], leaf: Optional[str]) -> Optional[str]:
-    ...
+def with_prefix(prefix: Optional[str], leaf: Optional[str]) -> Optional[str]: ...
 
 
 def with_prefix(prefix: Optional[str], leaf: Optional[str]) -> Optional[str]:

{haliax-1.4.dev409 → haliax-1.4.dev411}/src/haliax/_src/util.py

@@ -1,5 +1,9 @@
-from typing import Callable, MutableMapping, Sequence, TypeAlias, TypeVar
+# Copyright 2025 The Levanter Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
 
+from typing import Callable, MutableMapping, Sequence, TypeAlias, TypeVar
 
 T = TypeVar("T")
 U = TypeVar("U")

{haliax-1.4.dev409 → haliax-1.4.dev411}/src/haliax/axis.py

@@ -1,3 +1,8 @@
+# Copyright 2025 The Levanter Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+
 import typing
 from dataclasses import dataclass
 from math import prod
@@ -88,8 +93,7 @@ def selects_axis(selector: AxisSelection, selected: AxisSelection) -> bool:
     return True
 
 
-class _Sentinel:
-    ...
+class _Sentinel: ...
 
 
 def is_axis_compatible(ax1: AxisSelector, ax2: AxisSelector):
@@ -140,23 +144,19 @@ def axis_spec_to_shape_dict(axis_spec: AxisSelection) -> dict[str, Optional[int]
 
 
 @typing.overload
-def axis_spec_to_tuple(axis_spec: ShapeDict) -> tuple[Axis, ...]:
-    ...
+def axis_spec_to_tuple(axis_spec: ShapeDict) -> tuple[Axis, ...]: ...
 
 
 @typing.overload
-def axis_spec_to_tuple(axis_spec: AxisSpec) -> tuple[Axis, ...]:
-    ...
+def axis_spec_to_tuple(axis_spec: AxisSpec) -> tuple[Axis, ...]: ...
 
 
 @typing.overload
-def axis_spec_to_tuple(axis_spec: PartialShapeDict) -> tuple[AxisSelector, ...]:
-    ...
+def axis_spec_to_tuple(axis_spec: PartialShapeDict) -> tuple[AxisSelector, ...]: ...
 
 
 @typing.overload
-def axis_spec_to_tuple(axis_spec: AxisSelection) -> tuple[AxisSelector, ...]:
-    ...
+def axis_spec_to_tuple(axis_spec: AxisSelection) -> tuple[AxisSelector, ...]: ...
 
 
 def axis_spec_to_tuple(axis_spec: AxisSelection) -> tuple[AxisSelector, ...]:
@@ -230,23 +230,19 @@ def concat_axes(a1, a2):
 
 
 @typing.overload
-def union_axes(a1: ShapeDict, a2: AxisSpec) -> ShapeDict:
-    ...
+def union_axes(a1: ShapeDict, a2: AxisSpec) -> ShapeDict: ...
 
 
 @typing.overload
-def union_axes(a1: AxisSpec, a2: ShapeDict) -> ShapeDict:
-    ...
+def union_axes(a1: AxisSpec, a2: ShapeDict) -> ShapeDict: ...
 
 
 @typing.overload
-def union_axes(a1: AxisSpec, a2: AxisSpec) -> AxisSpec:
-    ...
+def union_axes(a1: AxisSpec, a2: AxisSpec) -> AxisSpec: ...
 
 
 @typing.overload
-def union_axes(a1: AxisSelection, a2: AxisSelection) -> AxisSelection:
-    ...
+def union_axes(a1: AxisSelection, a2: AxisSelection) -> AxisSelection: ...
 
 
 def union_axes(a1: AxisSelection, a2: AxisSelection) -> AxisSelection:
@@ -372,23 +368,19 @@ def without_axes(axis_spec: AxisSelection, to_remove: AxisSelection, allow_misma
 
 
 @typing.overload
-def unsize_axes(axis_spec: PartialShapeDict, to_unsize: AxisSelection) -> PartialShapeDict:
-    ...
+def unsize_axes(axis_spec: PartialShapeDict, to_unsize: AxisSelection) -> PartialShapeDict: ...
 
 
 @typing.overload
-def unsize_axes(axis_spec: AxisSelection, to_unsize: AxisSelection) -> AxisSelection:
-    ...
+def unsize_axes(axis_spec: AxisSelection, to_unsize: AxisSelection) -> AxisSelection: ...
 
 
 @typing.overload
-def unsize_axes(axis_spec: PartialShapeDict) -> PartialShapeDict:
-    ...
+def unsize_axes(axis_spec: PartialShapeDict) -> PartialShapeDict: ...
 
 
 @typing.overload
-def unsize_axes(axis_spec: AxisSelection) -> AxisSelection:
-    ...
+def unsize_axes(axis_spec: AxisSelection) -> AxisSelection: ...
 
 
 def unsize_axes(axis_spec: AxisSelection, to_unsize: Optional[AxisSelection] = None) -> AxisSelection:
@@ -424,13 +416,11 @@ def unsize_axes(axis_spec: AxisSelection, to_unsize: Optional[AxisSelection] = N
 
 
 @overload
-def replace_axis(axis_spec: AxisSpec, old: AxisSelector, new: AxisSpec) -> AxisSpec:
-    ...
+def replace_axis(axis_spec: AxisSpec, old: AxisSelector, new: AxisSpec) -> AxisSpec: ...
 
 
 @overload
-def replace_axis(axis_spec: AxisSelection, old: AxisSelector, new: AxisSelection) -> AxisSelection:
-    ...
+def replace_axis(axis_spec: AxisSelection, old: AxisSelector, new: AxisSelection) -> AxisSelection: ...
 
 
 def replace_axis(axis_spec: AxisSelection, old: AxisSelector, new: AxisSelection) -> AxisSelection:
@@ -466,13 +456,11 @@ def replace_axis(axis_spec: AxisSelection, old: AxisSelector, new: AxisSelection
 
 
 @overload
-def intersect_axes(ax1: ShapeDict, ax2: AxisSelection) -> ShapeDict:
-    ...
+def intersect_axes(ax1: ShapeDict, ax2: AxisSelection) -> ShapeDict: ...
 
 
 @overload
-def intersect_axes(ax1: tuple[AxisSelector, ...], ax2: AxisSpec) -> tuple[Axis, ...]:
-    ...
+def intersect_axes(ax1: tuple[AxisSelector, ...], ax2: AxisSpec) -> tuple[Axis, ...]: ...
 
 
 @overload
@@ -550,13 +538,11 @@ def axis_size(ax: AxisSpec) -> int:
 
 
 @typing.overload
-def resolve_axis(axis_spec: AxisSpec, axis_selection: AxisSelector) -> Axis:
-    ...
+def resolve_axis(axis_spec: AxisSpec, axis_selection: AxisSelector) -> Axis: ...
 
 
 @typing.overload
-def resolve_axis(axis_spec: AxisSpec, axis_selection: AxisSelection) -> AxisSpec:
-    ...
+def resolve_axis(axis_spec: AxisSpec, axis_selection: AxisSelection) -> AxisSpec: ...
 
 
 def resolve_axis(axis_spec: AxisSpec, axis_selection: AxisSelection) -> AxisSpec: