haliax 1.4.dev318__tar.gz → 1.4.dev321__tar.gz

{haliax-1.4.dev318 → haliax-1.4.dev321}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: haliax
-Version: 1.4.dev318
+Version: 1.4.dev321
 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.dev318 → haliax-1.4.dev321}/docs/api.md

@@ -31,10 +31,11 @@ Occasionally, an axis size can be inferred in some circumstances but not others.
 ::: haliax.axis.axis_name
 ::: haliax.axis.concat_axes
 ::: haliax.axis.union_axes
+::: haliax.axis.intersect_axes
 ::: haliax.axis.eliminate_axes
 ::: haliax.axis.without_axes
-::: haliax.axis.overlapping_axes
 ::: haliax.axis.selects_axis
+::: haliax.axis.overlapping_axes
 ::: haliax.axis.is_axis_compatible
 
 

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

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

{haliax-1.4.dev318 → haliax-1.4.dev321}/src/haliax/_src/einsum.py

@@ -121,10 +121,13 @@ def _output_only_named_einsum(equation, arrays, rhs, axis_aliases):
     used_aliases = set()
 
     input_axis_names = set(ax.name for ax in _all_input_axes(arrays))
+    has_ellipsis = False
 
     for capture in rhs.captures:
         if capture is Ellipsis:
-            raise_parse_error("Can't use ellipsis on the rhs of an einsum without an lhs", equation, None)
+            #     raise_parse_error("Can't use ellipsis on the rhs of an einsum without an lhs", equation, None)
+            out_axes.append(Ellipsis)
+            has_ellipsis = True
         elif capture.binding is None or len(capture.axes) > 1:
             raise_parse_error(
                 "Parenthesized axes are not currently supported in the output of an einsum",
@@ -150,7 +153,6 @@ def _output_only_named_einsum(equation, arrays, rhs, axis_aliases):
                     )
 
                 name = ax_name
-                used_axes.add(name)
 
             if name in out_axes:
                 raise_parse_error(
@@ -160,8 +162,18 @@ def _output_only_named_einsum(equation, arrays, rhs, axis_aliases):
             if name not in input_axis_names:
                 raise_parse_error(f"Axis {name} not found in any of the input arrays", equation, capture.char_range)
 
+            used_axes.add(name)
             out_axes.append(name)
 
+    # if there's an ellipsis, put all unused axes in the ellipsis
+    if has_ellipsis:
+        all_input_axes = _all_input_axes(arrays)
+        unmentioned = [ax.name for ax in all_input_axes if ax.name not in used_axes]
+        ellipsis_index = out_axes.index(Ellipsis)
+        out_axes = out_axes[:ellipsis_index] + unmentioned + out_axes[ellipsis_index + 1 :]
+
+        used_axes = set(out_axes)
+
     _check_for_unused_aliases(axis_aliases, used_aliases, equation)
 
     spec = _make_einsum_spec(arrays, out_axes)

{haliax-1.4.dev318 → haliax-1.4.dev321}/src/haliax/axis.py

@@ -297,21 +297,21 @@ def replace_axis(axis_spec: AxisSelection, old: AxisSelector, new: AxisSelection
 
 
 @overload
-def overlapping_axes(ax1: AxisSpec, ax2: AxisSelection) -> Tuple[Axis, ...]:
+def intersect_axes(ax1: AxisSpec, ax2: AxisSelection) -> Tuple[Axis, ...]:
     ...
 
 
 @overload
-def overlapping_axes(ax1: AxisSelection, ax2: AxisSpec) -> Tuple[Axis, ...]:
+def intersect_axes(ax1: AxisSelection, ax2: AxisSpec) -> Tuple[Axis, ...]:
     ...
 
 
 @overload
-def overlapping_axes(ax1: AxisSelection, ax2: AxisSelection) -> Tuple[AxisSelector, ...]:
+def intersect_axes(ax1: AxisSelection, ax2: AxisSelection) -> Tuple[AxisSelector, ...]:
     ...
 
 
-def overlapping_axes(ax1: AxisSelection, ax2: AxisSelection) -> Tuple[AxisSelector, ...]:
+def intersect_axes(ax1: AxisSelection, ax2: AxisSelection) -> Tuple[AxisSelector, ...]:
     """Returns a tuple of axes that are present in both ax1 and ax2.
     The returned order is the same as ax1.
     """
@@ -339,6 +339,23 @@ def overlapping_axes(ax1: AxisSelection, ax2: AxisSelection) -> Tuple[AxisSelect
     return tuple(out)
 
 
+def overlapping_axes(ax1: AxisSelection, ax2: AxisSelection) -> Tuple[str, ...]:
+    """
+    Like intersect_axes, but returns the names instead of the axes themselves.
+    Unlike intersect_axes, it does not throw an error if the sizes of a common axis are
+    different.
+
+    The returned order is the same as in ax1.
+    """
+    ax1 = ensure_tuple(ax1)
+    ax2 = ensure_tuple(ax2)
+    ax1_names = map(axis_name, ax1)
+    ax2_names = set(map(axis_name, ax2))
+
+    out = tuple(name for name in ax1_names if name in ax2_names)
+    return out
+
+
 @overload
 def axis_name(ax: AxisSelector) -> str:  # type: ignore
     ...
@@ -555,6 +572,7 @@ __all__ = [
     "dslice",
     "dblock",
     "eliminate_axes",
+    "intersect_axes",
     "is_axis_compatible",
     "overlapping_axes",
     "replace_axis",

{haliax-1.4.dev318 → haliax-1.4.dev321}/src/haliax/core.py

@@ -689,9 +689,9 @@ def take(array: NamedArray, axis: AxisSelector, index: Union[int, NamedArray]) -
         remaining_axes = eliminate_axes(array.axes, axis)
         # axis order is generally [array.axes[:axis_index], index.axes, array.axes[axis_index + 1 :]]
         # except that index.axes may overlap with array.axes
-        overlapping_axes: AxisSpec = haliax.axis.overlapping_axes(remaining_axes, index.axes)
+        intersecting_axes: AxisSpec = haliax.axis.intersect_axes(remaining_axes, index.axes)
 
-        if overlapping_axes:
+        if intersecting_axes:
             # if the eliminated axis is also in the index, we rename it to a dummy axis that we can broadcast over it
             need_to_use_dummy_axis = index._lookup_indices(axis.name) is not None
             if need_to_use_dummy_axis:

{haliax-1.4.dev318 → haliax-1.4.dev321}/tests/test_axis.py

@@ -1,6 +1,6 @@
 import pytest
 
-from haliax.axis import Axis, eliminate_axes, make_axes, rearrange_for_partial_order
+from haliax.axis import Axis, eliminate_axes, make_axes, overlapping_axes, rearrange_for_partial_order
 
 
 def test_eliminate_axes():
@@ -133,3 +133,17 @@ def test_duplicate_elements_errors():
 
     with pytest.raises(ValueError):
         rearrange_for_partial_order(partial_order, candidates)
+
+
+def test_overlapping_axes_with_different_sizes():
+    A1 = Axis("A", 10)
+    A2 = Axis("A", 12)
+    B = Axis("B", 14)
+    C = Axis("C", 16)
+    D = Axis("D", 18)
+
+    ax1 = (A1, B, C)
+    ax2 = (A2, C, D)
+
+    overlapping_names = overlapping_axes(ax1, ax2)  # Should not error
+    assert overlapping_names == ("A", "C")

{haliax-1.4.dev318 → haliax-1.4.dev321}/tests/test_einsum.py

@@ -250,9 +250,6 @@ def test_einsum_various_errors():
     m1 = hax.ones((Height, Hidth, Depth))
     m2 = hax.ones((Depth, Hidth, Height))
 
-    with pytest.raises(ValueError, match="Can't use ellipsis"):
-        einsum("-> ...", m1, m2)
-
     with pytest.raises(ValueError, match="multiple times"):
         einsum("-> Height Height", m1, m2)
 
@@ -330,6 +327,12 @@ def test_einsum_output_only_mode():
     assert jnp.all(jnp.equal(einsum("-> Height Width", m1, m2).array, jnp.einsum("ijk,kji->ij", m1.array, m2.array)))
     assert jnp.all(jnp.equal(einsum("-> Height", m1).array, jnp.einsum("ijk->i", m1.array)))
 
+    assert jnp.all(jnp.equal(einsum("-> ...", m1, m2).array, jnp.einsum("ijk,kji->ijk", m1.array, m2.array)))
+    assert jnp.all(jnp.equal(einsum("-> ... Width", m1, m2).array, jnp.einsum("ijk,kji->ikj", m1.array, m2.array)))
+    assert jnp.all(
+        jnp.equal(einsum("-> Depth ... Width", m1, m2).array, jnp.einsum("ijk,kji->kij", m1.array, m2.array))
+    )
+
     with pytest.raises(ValueError):
         einsum("-> Q Width", m1)
 

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

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