haliax 1.4.dev363__tar.gz → 1.4.dev364__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev363
+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.dev363 → 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.dev363 → 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.dev363 → 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,6 +39,7 @@ from .axis import (
     replace_axis,
     resolve_axis,
     selects_axis,
+    to_jax_shape,
 )
 from .core import (
     Named,
@@ -105,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:
@@ -155,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

{haliax-1.4.dev363 → 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(