haliax 1.4.dev363__tar.gz → 1.4.dev365__tar.gz

{haliax-1.4.dev363 → haliax-1.4.dev365}/.pre-commit-config.yaml

@@ -2,7 +2,7 @@
 # See https://pre-commit.com/hooks.html for more hooks
 exclude: ".git"
 default_stages:
-    - commit
+    - pre-commit
 fail_fast: true
 
 repos:
@@ -16,24 +16,24 @@ repos:
     -   id: check-merge-conflict
     -   id: check-added-large-files
 
--   repo: https://github.com/psf/black
-    rev: 22.3.0
-    hooks:
-    -   id: black
-
--   repo: https://github.com/timothycrosley/isort
-    rev: 5.11.5
-    hooks:
-    -   id: isort
+- 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/PyCQA/flake8
-    rev: 3.9.2
-    hooks:
-    -   id: flake8
-        additional_dependencies: [flake8-isort]
+#- repo: local
+#  hooks:
+#    - id: ty-check
+#      name: ty-check
+#      language: python
+#      entry: ty check
+#      pass_filenames: false
+#      args: [--python=.venv/]
+#      additional_dependencies: [ty]
 
 -   repo: https://github.com/pre-commit/mirrors-mypy
-    rev: 'v1.5.1'
+    rev: 'v1.16.1'
     hooks:
     -   id: mypy
         args: [--ignore-missing-imports, --check-untyped-defs]

haliax-1.4.dev365/AGENTS.md

@@ -0,0 +1,79 @@
+# Haliax LLM Agent Guidelines
+
+This document summarizes important conventions for contributing code or documentation to the Haliax
+repository. Follow these notes when implementing new features or fixing bugs.
+
+## General Guidelines
+
+* **Get better.** Whenever you discover something missing from these guidelines, or the requester
+  suggests a better way to do something, please update this document. The goal is to make it easier for
+  everyone to contribute and maintain the codebase. Generally speaking, you should add bullets or new sections.
+  Be sure to do this when directed to. For example, if directed that you should never relax tolerances in
+  floating point tests, add that to the list.
+* **Playbooks.** Sometimes, there are repeatable tasks (e.g. porting models) for which we follow a standard set of steps.
+  Please reference `.playbooks/` to see what playbooks are available, or see the list below. If you want to add a playbook
+  write a markdown doc named e.g. `.playbooks/port-models.md` and add a pointer to it in the list below.
+
+## Playbook
+
+- At the moment, there are no playbooks available. If you have a repeatable task that you think
+  should be documented, please create a new markdown file in `.playbooks/` and add it to the list above.
+
+## Code Style
+
+* **Python version**: the project targets Python >=3.10.
+* **Formatting and Linting**: We use `ruff` via `pre-commit`.
+* **Typing**: the code base uses `mypy` for static type checking. `mypy` is run by pre‑commit and the
+  configuration is found in `pyproject.toml`.
+* **Run `pre-commit run --all-files`** before committing. The CI workflows run the same checks.
+* **Doc Strings**: All public functions, classes, and modules should have docstrings, unless
+  their purpose is painfully obvious. Use
+  [Google style](https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings) for
+  consistency.
+* **Commenting**: Use comments to explain why something is done a certain way, especially if it is not
+  immediately obvious. Avoid commenting on every line of code; focus on the intent and purpose of
+  complex logic. Demarcating logical groups of code with comments is encouraged, unless it is better
+  to refactor the code into smaller functions or classes.
+* **Mkdocs**: We use [Mkdocs](https://www.mkdocs.org/) for documentation. The main documentation is in
+  the `docs` directory. Use Markdown for writing docs, and follow the existing structure. When linking to
+  symbols, prefer using mkdocs-style links (e.g. With a custom title: `[full.path.object2][]` or
+  `[Object 1][full.path.object1]`)
+* **Documentation**: When adding new features, ensure that the documentation is updated accordingly.
+  This includes updating the Mkdocs files and any relevant docstrings. If you add a new module or
+  significant functionality, consider adding a dedicated section in the documentation.
+
+## Testing
+
+* Tests are executed with `pytest`. The default workflow runs
+  `pytest tests -m "not entry and not slow and not ray"`.
+* In general, never relax tolerances in floating point tests unless specifically discussed with the
+  team. Use `assert_allclose` with appropriate tolerances for numerical comparisons. We typically use
+  1e-4 for more complex modules, and 1e-5 for simpler ones.
+* Tests should be reasonably fast. Mark long-running tests with @pytest.mark.slow so they are excluded from the default suite.
+* Always mark tests that depend on pytorch with `@skip_if_no_torch` to ensure they are skipped
+  when PyTorch is not available. This is particularly important for tests that require PyTorch-specific
+  functionality.
+
+
+## Design Preferences
+
+* **Generic code**: many utilities are written with Python generics and dataclasses. Where possible,
+  write reusable functions or classes that operate over TypeVars instead of hard coding concrete types.
+* **Configurations**: configuration files are dataclasses loaded via `draccus`. Keep configs
+  declarative and typed.
+* **Reproducibility**: Levanter aims for deterministic training where possible. Avoid sources of
+  nondeterminism unless explicitly required.
+* Prefer Stacked with fold or scan over writing custom loops, for better compile times and gradient checkpointing support
+
+## Library conventions
+- Haliax revolves around `NamedArray` and explicit `Axis` objects. Prefer APIs that accept
+  axes or axis names rather than hard‑coding positional dimensions.
+- Utilities should be written so they work with arbitrary axis names. Avoid relying on
+  fixed axis orders when possible.
+- Use the provided modules in `haliax.nn` or Equinox when building neural network layers.
+- Type annotations can use named shapes shorthand provided in `haliax.haxtyping`: `ht.f32[NamedArray, "batch"]`
+  for a float32 array with a "batch" axis, or `ht.Float[NamedArray, "batch"]` for any floating point dtype.
+
+## Documentation
+- Public functions and modules require docstrings. If behavior is non‑obvious,
+  add examples in `docs/`.

{haliax-1.4.dev363 → haliax-1.4.dev365}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev363
+Version: 1.4.dev365
 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/

{haliax-1.4.dev363 → haliax-1.4.dev365}/docs/api.md

@@ -35,7 +35,6 @@ Occasionally, an axis size can be inferred in some circumstances but not others.
 ::: haliax.axis.eliminate_axes
 ::: haliax.axis.without_axes
 ::: haliax.axis.selects_axis
-::: haliax.axis.overlapping_axes
 ::: haliax.axis.is_axis_compatible
 
 

{haliax-1.4.dev363 → haliax-1.4.dev365}/docs/nn.md

@@ -55,7 +55,6 @@ We don't provide an explicit attention module, but we do provide an attention fu
 
 :::haliax.nn.attention.dot_product_attention
 :::haliax.nn.attention.dot_product_attention_weights
-:::haliax.nn.attention.self_attention
 
 ### Masks
 ::: haliax.nn.attention.causal_mask

haliax-1.4.dev365/docs/typing.md

@@ -0,0 +1,100 @@
+from haliax import NamedArrayfrom haliax import NamedArray
+
+# NamedArray Type Annotations
+
+Haliax supports an extension to [`jaxtyping`](https://docs.kidger.site/jaxtyping/)
+that allows you to annotate functions and methods that take or return
+[`NamedArray`][haliax.core.NamedArray] objects. If you are familiar with
+[`jaxtyping`](https://docs.kidger.site/jaxtyping/), the syntax is very similar.
+In fact, for non-NamedArrays, it is exactly the same.
+
+```python
+from haliax import NamedArray
+import haliax.haxtyping as ht
+
+def foo(x: ht.Float[NamedArray, "batch embed ..."]):
+    ...
+```
+
+At runtime you can verify that a `NamedArray` conforms to a particular
+annotation using `matches_axes`:
+
+```python
+if not arr.matches_axes(Float[NamedArray, "batch embed ..."]):
+    raise ValueError("unexpected axes")
+```
+
+## DType-aware annotations
+
+Sometimes it is useful to express both the axes **and** the dtype in the type
+annotation.  The :mod:`haliax.typing` module defines symbolic types for all of
+JAX's common dtypes that can be indexed just like ``Named``.  In documentation
+examples we'll use ``import haliax.typing as ht``:
+
+```python
+import haliax.haxtyping as ht
+
+def foo(x: ht.f32[NamedArray, "batch"]):
+    ...
+
+def bar(x: ht.i32[NamedArray, "batch"]):
+    ...
+```
+
+For convenience the module also provides aggregate categories ``Float``,
+``Complex``, ``Int`` and ``UInt`` that match any floating point, complex,
+signed integer or unsigned integer dtype respectively:
+
+```python
+def baz(x: ht.Float[NamedArray, "batch"]):
+    ...
+```
+
+At runtime ``matches_axes`` also checks the dtype when one is present:
+
+```python
+from haliax import Axis, zeros
+import haliax.haxtyping as ht
+
+arr = zeros({"batch": 4})
+assert arr.matches_axes(ht.f32["batch"])  # dtype and axes both match
+```
+
+## FAQ
+
+### Why not use `NamedArray` directly in type annotations?
+
+Using `NamedArray` directly in type annotations doesn't work well with
+type checkers like `mypy` or `pyright`. These tools expect types to be
+subscripted with other types or forward references (which are strings).
+Using `NamedArray` directly would lead to type errors.
+
+### Why not use `jaxtyping` directly?
+
+While `jaxtyping` is a powerful library for type annotations in JAX, it does not
+support `NamedArray` objects directly. The `haliax.haxtyping` module extends
+`jaxtyping` to include `NamedArray` support, allowing you to annotate functions
+and methods that take or return `NamedArray` objects with specific axes and dtypes.
+
+### Why do I have to specify the `NamedArray` type in the annotation?
+
+I hate this, but it's the only way to get type checkers like `mypy` and `pyright` to understand that the type is
+a `NamedArray`. Underneath the hood, during type checking, `jaxtyping.Float` (and `haxtyping.Float`) are
+essentially type aliases of [`Annotated`](https://docs.python.org/3/library/typing.html#typing.Annotated)
+with the `NamedArray` type. There's no other way I could find to get type checkers to understand that the type is a
+`NamedArray` or to accept strings like `"batch embed ..."` as valid type annotations.
+
+### How do I use single axes in type annotations with flake or ruff.
+
+Like `jaxtyping`, you need to prepend a space before the axis name to use single axes in type annotations with
+flake or ruff. For example, to use a single axis named `batch`, you would write:
+
+```python
+def foo(x: ht.Float[NamedArray, " batch"]):
+    ...
+```
+
+Then suppress F722 in your linter to suppress that error.
+
+See the [jaxtyping documentation](https://docs.kidger.site/jaxtyping/faq/#flake8-or-ruff-are-throwing-an-error) for more
+details on the workaround.

{haliax-1.4.dev363 → haliax-1.4.dev365}/pyproject.toml

@@ -71,3 +71,24 @@ src_paths = ["src", "tests"]
 "Homepage" = "https://github.com/stanford-crfm/haliax"
 "Bug Tracker" = "https://github.com/stanford-crfm/haliax/issues/"
 "Documentation" = "https://haliax.readthedocs.io/en/latest/"
+
+[tool.ruff.lint]
+ignore = [ "E203", "E501", "W605", "F821", "E266", "F722", "E731", "E741" ]
+
+[tool.setuptools.package-data]
+"haliax" = ["*.pyi"]
+
+[tool.coverage.report]
+exclude_also = [
+    "def __repr__",
+    "if self.debug:",
+    "if settings.DEBUG",
+    "raise AssertionError",
+    "raise NotImplementedError",
+    "if 0:",
+    "if __name__ == .__main__.:",
+    "if TYPE_CHECKING:",
+    "class .*\\bProtocol\\):",
+    "@(abc\\.)?abstractmethod",
+    "[.][.][.]"
+    ]

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

@@ -0,0 +1 @@
+__version__ = "1.4.dev365"

{haliax-1.4.dev363 → haliax-1.4.dev365}/src/haliax/__init__.py

@@ -1,4 +1,4 @@
-import typing
+import typing as t
 from typing import Optional, Sequence
 
 import jax
@@ -29,6 +29,7 @@ from .axis import (
     AxisSpec,
     axis_name,
     axis_size,
+    axis_spec_to_tuple,
     concat_axes,
     dblock,
     ds,
@@ -38,13 +39,11 @@ from .axis import (
     replace_axis,
     resolve_axis,
     selects_axis,
+    to_jax_shape,
 )
 from .core import (
-    Named,
     NamedArray,
-    NamedArrayAxes,
-    NamedArrayAxesSpec,
-    NamedOrNumeric,
+    NamedArrayAxes, NamedArrayAxesSpec, NamedOrNumeric,
     are_shape_checks_enabled,
     broadcast_arrays,
     broadcast_axis,
@@ -64,6 +63,7 @@ from .core import (
     unflatten_axis,
     updated_slice,
 )
+from .haxtyping import Named
 from .hof import fold, map, scan, vmap
 from .jax_utils import tree_checkpoint_name
 from .ops import clip, isclose, pad_left, trace, tril, triu, where
@@ -81,8 +81,8 @@ from .wrap import (
 )
 
 
-T = typing.TypeVar("T")
-A = typing.TypeVar("A", Scalar, NamedArray, jnp.ndarray)
+T = t.TypeVar("T")
+A = t.TypeVar("A", Scalar, NamedArray, jnp.ndarray)
 
 
 # creation routines
@@ -105,8 +105,8 @@ def full(shape: AxisSpec, fill_value: T, dtype: Optional[DTypeLike] = None) -> N
     if isinstance(shape, Axis):
         return NamedArray(jnp.full(shape=shape.size, fill_value=fill_value, dtype=dtype), (shape,))
     else:
-        x_shape = tuple(x.size for x in shape)
-        return NamedArray(jnp.full(shape=x_shape, fill_value=fill_value, dtype=dtype), tuple(shape))
+        x_shape = to_jax_shape(shape)
+        return NamedArray(jnp.full(shape=x_shape, fill_value=fill_value, dtype=dtype), shape)
 
 
 def zeros_like(a: NamedArray, dtype=None) -> NamedArray:
@@ -146,16 +146,11 @@ def arange(axis: AxisSpec, *, start=0, step=1, dtype: Optional[DTypeLike] = None
     ```
 
     """
-    from haliax.jax_utils import to_jax_shape
-    from haliax.util import ensure_tuple
-
-    # if start is a tracer, we need to be a bit cleverer since arange doesn't support tracers
-    # return NamedArray(jnp.arange(start, stop, step, dtype=dtype), (axis,))
     size = axis_size(axis)
 
     arr = jax.lax.iota(dtype=dtype or jnp.result_type(start), size=size) * step + start
     arr = arr.reshape(to_jax_shape(axis))
-    return NamedArray(arr, ensure_tuple(axis))
+    return NamedArray(arr, axis_spec_to_tuple(axis))
 
 
 # TODO: add overrides for arraylike start/stop to linspace, logspace, geomspace
@@ -272,7 +267,7 @@ def concatenate(axis: AxisSelector, arrays: Sequence[NamedArray]) -> NamedArray:
     if axis_index is None:
         raise ValueError(f"Axis {aname} not found in 0th array {arrays[0]}")
 
-    axes: typing.Tuple[AxisSelector, ...] = arrays[0].axes
+    axes: tuple[AxisSelector, ...] = arrays[0].axes
     # we want to use the axis name for `axis`, because it's not uncommon for those to be different lengths in the arrays
     axes = axes[:axis_index] + (aname,) + axes[axis_index + 1 :]
     arrays = [a.rearrange(axes) for a in arrays]
@@ -929,9 +924,6 @@ __all__ = [
     "make_axes",
     "axis_name",
     "axis_size",
-    "NamedArrayAxesSpec",
-    "NamedArrayAxes",
-    "Named",
     "NamedArray",
     "broadcast_to",
     "broadcast_axis",
@@ -1104,4 +1096,14 @@ __all__ = [
     "is_named_array",
     "tree_checkpoint_name",
     "ScanCheckpointPolicy",
+    "quantization",
+    "util",
+    "einsum",
+    "broadcast_arrays",
+    "unflatten_axis",
+    "ReductionFunction",
+    "SimpleReductionFunction",
+    "NamedArrayAxes",
+    "NamedArrayAxesSpec",
+    "Named",
 ]

{haliax-1.4.dev363 → haliax-1.4.dev365}/src/haliax/_src/dot.py

@@ -4,7 +4,6 @@ import warnings
 from typing import Dict, Optional, Tuple
 
 import jax
-import jax.numpy as jnp
 
 import haliax
 from haliax.axis import (
@@ -12,6 +11,7 @@ from haliax.axis import (
     AxisSelection,
     PartialAxisSpec,
     axis_name,
+    axis_spec_to_shape_dict,
     eliminate_axes,
     rearrange_for_partial_order,
     union_axes,
@@ -19,7 +19,6 @@ from haliax.axis import (
 from haliax.core import NamedArray
 from haliax.jax_utils import _jittable_dg_einsum
 from haliax.types import DTypeLike, PrecisionLike
-from haliax.util import ensure_tuple
 
 
 # deprecated overload
@@ -140,8 +139,8 @@ def dot(
     if axis is None:
         jax_str = f"contract {', '.join(axis_name(ax) for ax in all_axes)} -> <scalar>"
     else:
-        axis = ensure_tuple(axis)
-        jax_str = f"contract {', '.join(axis_name(ax) for ax in axis)} -> {', '.join(a.name for a in output_axes)}"
+        axis = axis_spec_to_shape_dict(axis)
+        jax_str = f"contract {', '.join(axis)} -> {', '.join(a.name for a in output_axes)}"
 
     with jax.named_scope(jax_str):
         output = _jittable_dg_einsum(

{haliax-1.4.dev363 → haliax-1.4.dev365}/src/haliax/_src/fp8.py

@@ -1,4 +1,3 @@
-import warnings
 from functools import partial
 
 from jax import custom_jvp, custom_vjp, lax

{haliax-1.4.dev363 → haliax-1.4.dev365}/src/haliax/_src/state_dict.py

@@ -15,8 +15,7 @@ from jaxtyping import PyTree
 
 import haliax.partitioning as partitioning
 from haliax._src.util import index_where
-from haliax.axis import Axis
-from haliax.core import NamedArray, flatten_axes, named
+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
 
@@ -209,14 +208,14 @@ def from_state_dict(tree: T, state_dict: StateDict, prefix: Optional[str] = None
             array = named(array, tree.axes)
             array = partitioning.auto_sharded(array)
 
-        return array
+        return array  # type: ignore
     elif is_jax_array_like(tree):
         if prefix is None:
             raise ValueError("Cannot extract a leaf value from a state dict without a prefix")
         # TODO: add "strict" flag so we can return None in cases where it's just missing
         return jnp.array(state_dict[prefix])
     elif tree is None:
-        return None
+        return None  # type: ignore
     else:
         if prefix is None:
             return tree