haliax 1.4.dev360__tar.gz → 1.4.dev363__tar.gz

{haliax-1.4.dev360 → haliax-1.4.dev363}/.github/workflows/run_tests.yaml

@@ -17,7 +17,7 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           pip install flake8 pytest
-          pip install jax==0.4.35 jaxlib==0.4.35 .[dev]
+          pip install -e .[dev]
       - name: Test with pytest
         run: |
           XLA_FLAGS=--xla_force_host_platform_device_count=8 PYTHONPATH=tests:src:. pytest tests

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

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

@@ -0,0 +1,61 @@
+# NamedArray Type Annotations
+
+Haliax supports a lightweight syntax for specifying the axes of a `NamedArray`
+in type annotations.  Internally, `Named[...]` expands to
+`Annotated[NamedArray, axes]`, so it works well with static type checkers like
+``mypy``.  The syntax mirrors normal indexing with axis names.  Some examples:
+
+```python
+from haliax import Named
+
+arr: Named["batch", "embed"]
+arr: Named["batch embed ..."]      # starts with these axes
+arr: Named["... embed"]             # ends with this axis
+arr: Named["batch ... embed"]       # contains these axes in order
+arr: Named[{"batch", "embed"}]      # has exactly these axes, order ignored
+arr: Named[{"batch", "embed", ...}] # has at least these axes
+```
+
+At runtime you can verify that a `NamedArray` conforms to a particular
+annotation using `matches_axes`:
+
+```python
+if not arr.matches_axes(Named["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.typing as ht
+
+def foo(x: ht.f32["batch"]):
+    ...
+
+def bar(x: ht.i32["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["batch"]):
+    ...
+```
+
+At runtime ``matches_axes`` also checks the dtype when one is present:
+
+```python
+from haliax import Axis, zeros
+import haliax.typing as ht
+
+arr = zeros({"batch": 4})
+assert arr.matches_axes(ht.f32["batch"])  # dtype and axes both match
+```

{haliax-1.4.dev360 → haliax-1.4.dev363}/mkdocs.yml

@@ -90,6 +90,7 @@ nav:
     - Indexing and Slicing: 'indexing.md'
     - Rearrange: 'rearrange.md'
     - Matrix Multiplication: 'matmul.md'
+    - Type Annotations: 'typing.md'
   - Higher Order Functions:
     - Scan and Fold: 'scan.md'
     - Vectorization: 'vmap.md'

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

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

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

@@ -40,7 +40,10 @@ from .axis import (
     selects_axis,
 )
 from .core import (
+    Named,
     NamedArray,
+    NamedArrayAxes,
+    NamedArrayAxesSpec,
     NamedOrNumeric,
     are_shape_checks_enabled,
     broadcast_arrays,
@@ -926,6 +929,9 @@ __all__ = [
     "make_axes",
     "axis_name",
     "axis_size",
+    "NamedArrayAxesSpec",
+    "NamedArrayAxes",
+    "Named",
     "NamedArray",
     "broadcast_to",
     "broadcast_axis",

{haliax-1.4.dev360 → haliax-1.4.dev363}/src/haliax/core.py

@@ -58,9 +58,139 @@ def are_shape_checks_enabled():
     return _ENABLE_SHAPE_CHECKS
 
 
+@dataclass(frozen=True)
+class NamedArrayAxes:
+    """Representation of a :class:`NamedArray`'s axes for type annotations."""
+
+    before: Tuple[str, ...]
+    """Names that must appear before any optional ellipsis."""
+
+    after: Tuple[str, ...] = ()
+    """Names that must appear after any optional ellipsis."""
+
+    ordered: bool = True
+    """Whether the axes must appear in this order."""
+
+    subset: bool = False
+    """If ``True``, other axes may appear where the ellipsis is located."""
+
+    dtype: typing.Any | None = None
+    """Optional dtype that the array should have."""
+
+    def __repr__(self) -> str:
+        dtype_prefix = ""
+        if self.dtype is not None:
+            dtype_obj = self.dtype
+            if hasattr(dtype_obj, "category"):
+                dtype_name = dtype_obj.name
+            else:
+                dtype_name = getattr(dtype_obj, "name", str(dtype_obj))
+            dtype_prefix = f"{dtype_name} "
+
+        if self.ordered:
+            parts = list(self.before)
+            if self.subset:
+                parts.append("...")
+            parts.extend(self.after)
+            spec = " ".join(parts)
+            return f"NamedArray[{dtype_prefix}{spec}]"
+        else:
+            part = ", ".join(self.before)
+            if self.subset:
+                if part:
+                    part += ", ..."
+                else:
+                    part = "..."
+            return f"NamedArray[{dtype_prefix}{{{part}}}]"
+
+
+# a specification for NamedArray axes used in type annotations
+NamedArrayAxesSpec = Union[
+    NamedArrayAxes,
+    str,
+    Sequence[str | EllipsisType],
+    set[str | EllipsisType],
+]
+
+
+def _parse_namedarray_axes(
+    item: NamedArrayAxesSpec | typing.Annotated["NamedArray", NamedArrayAxes]
+) -> NamedArrayAxes:
+    origin = typing.get_origin(item)
+    if origin is typing.Annotated:
+        args = typing.get_args(item)
+        if len(args) >= 2:
+            item = args[1]
+    if isinstance(item, NamedArrayAxes):
+        return item
+    if isinstance(item, str):
+        parts = item.split()
+        if parts.count("...") > 1:
+            raise TypeError("Only one ellipsis allowed in NamedArray typing spec")
+        if "..." in parts:
+            idx = parts.index("...")
+            before_parts = tuple(parts[:idx])
+            after_parts = tuple(parts[idx + 1 :])
+            return NamedArrayAxes(before_parts, after_parts, ordered=True, subset=True)
+        else:
+            return NamedArrayAxes(tuple(parts), (), ordered=True, subset=False)
+    if isinstance(item, set) or isinstance(item, frozenset):
+        subset = False
+        names_list: List[str] = []
+        for part in item:
+            if part is Ellipsis:
+                if subset:
+                    raise TypeError("Only one ellipsis allowed in NamedArray typing spec")
+                subset = True
+            else:
+                if not isinstance(part, str):
+                    raise TypeError(f"Invalid axis spec: {part}")
+                names_list.append(part)
+        return NamedArrayAxes(tuple(names_list), (), ordered=False, subset=subset)
+    if isinstance(item, (tuple, list)):
+        subset = False
+        before_list: List[str] = []
+        after_list: List[str] = []
+        cur_list = before_list
+        for part in item:
+            if part is Ellipsis:
+                if subset:
+                    raise TypeError("Only one ellipsis allowed in NamedArray typing spec")
+                subset = True
+                cur_list = after_list
+            else:
+                if not isinstance(part, str):
+                    raise TypeError(f"Invalid axis spec: {part}")
+                cur_list.append(part)
+        if subset:
+            return NamedArrayAxes(tuple(before_list), tuple(after_list), ordered=True, subset=True)
+        else:
+            return NamedArrayAxes(tuple(before_list), (), ordered=True, subset=False)
+    raise TypeError(f"Invalid NamedArray typing spec: {item}")
+
+
+class Named:
+    """Type annotation helper for :class:`NamedArray`.
+
+    ``Named["batch embed"]`` expands to ``Annotated[NamedArray, axes]`` so that
+    type checkers treat it as a ``NamedArray`` at static time while the axis
+    metadata is available at runtime via :func:`typing.get_args`.
+    """
+
+    def __class_getitem__(cls, item: NamedArrayAxesSpec) -> typing.Any:
+        axes = _parse_namedarray_axes(item)
+        return typing.Annotated[NamedArray, axes]
+
+
+class NamedArrayMeta(type):
+    def __getitem__(cls, item: NamedArrayAxesSpec) -> typing.Any:
+        axes = _parse_namedarray_axes(item)
+        return typing.Annotated[NamedArray, axes]
+
+
 @jax.tree_util.register_pytree_node_class
 @dataclass(frozen=True)
-class NamedArray:
+class NamedArray(metaclass=NamedArrayMeta):
     array: jnp.ndarray
     axes: Tuple[Axis, ...]
 
@@ -145,6 +275,45 @@ class NamedArray:
         """Returns true if the given axis is present in this NamedArray."""
         return self.axis_indices(axis) is not None
 
+    def matches_axes(self, spec: NamedArrayAxesSpec) -> bool:
+        """Check whether this NamedArray conforms to the given `NamedArray` type.
+
+        Parameters
+        ----------
+        spec : NamedArrayAxesSpec
+            The specification to check against. It can be produced via the
+            ``NamedArray[...]`` syntax or passed directly as a string or
+            sequence of axis names.
+        """
+
+        ann = _parse_namedarray_axes(spec)
+        if ann.dtype is not None:
+            dtype_spec = ann.dtype
+            if hasattr(dtype_spec, "category"):
+                if not jnp.issubdtype(self.dtype, dtype_spec.category):
+                    return False
+            elif self.dtype != dtype_spec:
+                return False
+
+        names = tuple(ax.name for ax in self.axes)
+        if ann.ordered:
+            if not ann.subset:
+                return names == ann.before
+            if len(names) < len(ann.before) + len(ann.after):
+                return False
+            if names[: len(ann.before)] != ann.before:
+                return False
+            if ann.after and names[-len(ann.after) :] != ann.after:
+                return False
+            return True
+        else:
+            name_set = set(names)
+            spec_set = set(ann.before)
+            if ann.subset:
+                return spec_set.issubset(name_set)
+            else:
+                return name_set == spec_set
+
     @overload
     def axis_size(self, axis: AxisSelector) -> int:  # type: ignore
         ...
@@ -1885,6 +2054,9 @@ def _convert_index_expr_to_dict(idx) -> dict[AxisSelector, NamedIndex]:
 
 
 __all__ = [
+    "NamedArrayAxesSpec",
+    "NamedArrayAxes",
+    "Named",
     "NamedArray",
     "named",
     "slice",

{haliax-1.4.dev360 → haliax-1.4.dev363}/src/haliax/nn/activations.py

@@ -89,7 +89,6 @@ def quick_gelu(x):
     return x * sigmoid(1.702 * x)
 
 
-
 def relu_squared(x: A) -> A:
     """ReLU squared activation function. jnp.square(jnp.maximum(0, x))"""
 

haliax-1.4.dev363/src/haliax/typing.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+import typing as tp
+from dataclasses import dataclass, replace
+
+import jax.numpy as jnp
+
+from .core import NamedArray, NamedArrayAxes, _parse_namedarray_axes
+
+
+@dataclass(frozen=True)
+class DTypeCategory:
+    """Represents a dtype category such as ``float`` or ``int``."""
+
+    name: str
+    category: tp.Any
+
+    def __repr__(self) -> str:  # pragma: no cover - trivial
+        return self.name
+
+
+def _wrap_namedarray_with_dtype(dtype):
+    class DTypeType:
+        def __class_getitem__(cls, axes_spec):
+            axes = _parse_namedarray_axes(axes_spec)
+            axes_with_dtype = replace(axes, dtype=dtype)
+            return tp.Annotated[NamedArray, axes_with_dtype]
+
+    return DTypeType
+
+
+def _wrap_namedarray_with_category(category: DTypeCategory):
+    class DTypeType:
+        def __class_getitem__(cls, axes_spec):
+            axes = _parse_namedarray_axes(axes_spec)
+            axes_with_dtype = replace(axes, dtype=category)
+            return tp.Annotated[NamedArray, axes_with_dtype]
+
+    return DTypeType
+
+
+f32 = _wrap_namedarray_with_dtype(jnp.float32)
+f64 = _wrap_namedarray_with_dtype(jnp.float64)
+f16 = _wrap_namedarray_with_dtype(jnp.float16)
+bf16 = _wrap_namedarray_with_dtype(jnp.bfloat16)
+
+i8 = _wrap_namedarray_with_dtype(jnp.int8)
+i16 = _wrap_namedarray_with_dtype(jnp.int16)
+i32 = _wrap_namedarray_with_dtype(jnp.int32)
+i64 = _wrap_namedarray_with_dtype(jnp.int64)
+
+u8 = _wrap_namedarray_with_dtype(jnp.uint8)
+u16 = _wrap_namedarray_with_dtype(jnp.uint16)
+u32 = _wrap_namedarray_with_dtype(jnp.uint32)
+u64 = _wrap_namedarray_with_dtype(jnp.uint64)
+
+bool_ = _wrap_namedarray_with_dtype(jnp.bool_)
+complex64 = _wrap_namedarray_with_dtype(jnp.complex64)
+complex128 = _wrap_namedarray_with_dtype(jnp.complex128)
+
+
+Float = _wrap_namedarray_with_category(DTypeCategory("float", jnp.floating))
+Complex = _wrap_namedarray_with_category(DTypeCategory("complex", jnp.complexfloating))
+Int = _wrap_namedarray_with_category(DTypeCategory("int", jnp.signedinteger))
+UInt = _wrap_namedarray_with_category(DTypeCategory("uint", jnp.unsignedinteger))
+
+
+__all__ = [
+    "f32",
+    "f64",
+    "f16",
+    "bf16",
+    "i8",
+    "i16",
+    "i32",
+    "i64",
+    "u8",
+    "u16",
+    "u32",
+    "u64",
+    "bool_",
+    "complex64",
+    "complex128",
+    "Float",
+    "Complex",
+    "Int",
+    "UInt",
+]

haliax-1.4.dev363/tests/test_dtype_typing.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+import typing
+
+import jax.numpy as jnp
+
+from haliax import Axis, NamedArray
+from haliax.typing import Float, Int, f32, i32
+
+
+def test_dtype_and_axes_annotation():
+    def foo(x: f32["batch embed"]):  # type: ignore  # noqa: F722
+        pass
+
+    ann = typing.get_args(typing.get_type_hints(foo, include_extras=True)["x"])
+    assert ann[0] is NamedArray
+    spec = ann[1]
+    assert spec.dtype == jnp.float32
+    assert spec.before == ("batch", "embed")
+
+
+def test_other_dtype_annotation():
+    def bar(x: i32["batch"]):  # type: ignore  # noqa: F722
+        pass
+
+    spec = typing.get_args(typing.get_type_hints(bar, include_extras=True)["x"])[1]
+    assert spec.dtype == jnp.int32
+    assert spec.before == ("batch",)
+
+
+def test_dtype_category_annotation_and_check():
+    def baz(x: Float["b"]):  # type: ignore  # noqa: F722
+        pass
+
+    spec = typing.get_args(typing.get_type_hints(baz, include_extras=True)["x"])[1]
+    assert str(spec.dtype) == "float"
+
+    B = Axis("b", 1)
+    arr = NamedArray(jnp.ones((B.size,), dtype=jnp.float32), (B,))
+    assert arr.matches_axes(Float["b"])  # type: ignore
+    assert not arr.matches_axes(Int["b"])  # type: ignore

haliax-1.4.dev363/tests/test_namedarray_typing.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+import typing
+
+import jax.numpy as jnp
+
+from haliax import Axis, Named, NamedArray
+from haliax.typing import Float, Int, f32, i32
+
+
+def test_namedarray_type_syntax():
+    axes1 = typing.get_args(NamedArray["batch", "embed"])[1]
+    axes2 = typing.get_args(Named["batch embed"])[1]  # type: ignore
+    assert axes1 == axes2
+
+    axes3 = typing.get_args(NamedArray["batch embed ..."])[1]
+    assert axes3.before == ("batch", "embed") and axes3.subset and axes3.after == ()
+
+    axes4 = typing.get_args(NamedArray[{"batch", "embed"}])[1]
+    assert set(axes4.before) == {"batch", "embed"} and not axes4.ordered
+
+    axes5 = typing.get_args(NamedArray[{"batch", "embed", ...}])[1]
+    assert set(axes5.before) == {"batch", "embed"} and not axes5.ordered and axes5.subset
+
+    axes6 = typing.get_args(NamedArray["... embed"])[1]
+    assert axes6.before == () and axes6.after == ("embed",) and axes6.subset
+
+    axes7 = typing.get_args(NamedArray["batch ... embed"])[1]
+    assert axes7.before == ("batch",) and axes7.after == ("embed",) and axes7.subset
+
+
+def test_named_param_annotation():
+    def foo(x: Named["batch", "embed"]):  # type: ignore
+        pass
+
+    axes = typing.get_args(typing.get_type_hints(foo, include_extras=True)["x"])[1]
+    assert axes.before == ("batch", "embed")
+
+
+def test_namedarray_runtime_check():
+    Batch = Axis("batch", 2)
+    Embed = Axis("embed", 3)
+    arr = NamedArray(jnp.zeros((Batch.size, Embed.size)), (Batch, Embed))
+    assert arr.matches_axes(NamedArray["batch", "embed"])
+    assert arr.matches_axes(Named["batch embed"])  # type: ignore
+    assert arr.matches_axes(NamedArray["batch embed ..."])
+    assert arr.matches_axes(NamedArray[{"batch", "embed"}])
+    assert arr.matches_axes(NamedArray[{"batch", "embed", ...}])
+    assert not arr.matches_axes(NamedArray["embed batch"])
+    assert not arr.matches_axes(NamedArray[{"batch", "foo", ...}])
+
+
+def test_namedarray_runtime_check_with_dtype():
+    Batch = Axis("batch", 2)
+    arr = NamedArray(jnp.zeros((Batch.size,), dtype=jnp.float32), (Batch,))
+    assert arr.matches_axes(f32["batch"])  # type: ignore
+    assert not arr.matches_axes(i32["batch"])  # type: ignore
+
+
+def test_namedarray_runtime_check_with_category():
+    B = Axis("batch", 1)
+    arr = NamedArray(jnp.zeros((B.size,), dtype=jnp.float32), (B,))
+    assert arr.matches_axes(Float["batch"])  # type: ignore
+    assert not arr.matches_axes(Int["batch"])  # type: ignore

{haliax-1.4.dev360 → haliax-1.4.dev363}/tests/test_nn.py

@@ -198,4 +198,4 @@ def test_relu_squared_scalar(use_jit):
     x_neg = -5.0
     expected_neg = 0.0
     actual_neg = f(x_neg)
-    assert jnp.allclose(actual_neg, expected_neg)
+    assert jnp.allclose(actual_neg, expected_neg)

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

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