haliax 1.4.dev319__tar.gz → 1.4.dev321__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: haliax
-Version: 1.4.dev319
+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.dev321/src/haliax/__about__.py

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

{haliax-1.4.dev319 → 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.dev319 → 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.dev319/src/haliax/__about__.py

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