haliax 1.4.dev362__tar.gz → 1.4.dev364__tar.gz

{haliax-1.4.dev362 → haliax-1.4.dev364}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev362
+Version: 1.4.dev364
 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.dev362 → haliax-1.4.dev364}/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.dev364/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.dev362 → haliax-1.4.dev364}/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.dev362 → haliax-1.4.dev364}/pyproject.toml

@@ -71,3 +71,19 @@ 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.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.dev364/src/haliax/__about__.py

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

{haliax-1.4.dev362 → haliax-1.4.dev364}/src/haliax/__init__.py

@@ -29,6 +29,7 @@ from .axis import (
     AxisSpec,
     axis_name,
     axis_size,
+    axis_spec_to_tuple,
     concat_axes,
     dblock,
     ds,
@@ -38,9 +39,13 @@ from .axis import (
     replace_axis,
     resolve_axis,
     selects_axis,
+    to_jax_shape,
 )
 from .core import (
+    Named,
     NamedArray,
+    NamedArrayAxes,
+    NamedArrayAxesSpec,
     NamedOrNumeric,
     are_shape_checks_enabled,
     broadcast_arrays,
@@ -102,8 +107,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:
@@ -152,7 +157,7 @@ def arange(axis: AxisSpec, *, start=0, step=1, dtype: Optional[DTypeLike] = None
 
     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
@@ -926,6 +931,9 @@ __all__ = [
     "make_axes",
     "axis_name",
     "axis_size",
+    "NamedArrayAxesSpec",
+    "NamedArrayAxes",
+    "Named",
     "NamedArray",
     "broadcast_to",
     "broadcast_axis",

{haliax-1.4.dev362 → haliax-1.4.dev364}/src/haliax/_src/dot.py

@@ -12,6 +12,7 @@ from haliax.axis import (
     AxisSelection,
     PartialAxisSpec,
     axis_name,
+    axis_spec_to_shape_dict,
     eliminate_axes,
     rearrange_for_partial_order,
     union_axes,
@@ -140,8 +141,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(