Repository: mrkwjc/einsumt

Issue:
`einsumt` fails if the shape of (ellipsis) trailing axes are not equal
...however, `np.einsum` does automatic broadcasting. Any ideas on implementing that?

Minimal working example:

```python
import numpy as np
from einsumt import einsumt

a = np.ones((3, 3, 1, 1))
b = np.ones((3, 3, 8, 100))

# this works
np.einsum("ij...,ij...->...", a, b)

# this fails
einsumt("ij...,ij...->...", a, b)

# ...this also fails
einsumt("ijmn,ijmn->mn", a, b)
```

Thanks for providing `einsumt`, it is a fantastic package. A little side-node: Mostly, `einsumt` is equal in execution times compared to just-in-time compiling the code (Numba). I'm using `einsumt` in my finite elements package [FElupe](https://github.com/adtzlr/felupe).

Best regards,
Andreas

Discussion:

Patch:
diff --git a/einsumt.py b/einsumt.py
index 85480c3..c43ff3d 100644
--- a/einsumt.py
+++ b/einsumt.py
@@ -123,8 +123,9 @@ def einsumt(*operands, **kwargs):
         for j in range(len(ops)):
             oj = ops[j]
             if cpos[j] >= 0:
-                jslice = cpos_slice[j] + (islice,)
-                oj = oj[jslice]
+                if oj.shape[cpos[j]] > 1:
+                    jslice = cpos_slice[j] + (islice,)
+                    oj = oj[jslice]
             args = args + (oj,)
         res += [pool.apply_async(np.einsum, args=args, kwds=kwargs)]
     res = [r.get() for r in res]
diff --git a/setup.py b/setup.py
index e3312c9..2584da6 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@ with open("README.md", "r") as fh:
 
 setuptools.setup(
     name="einsumt",
-    version="0.9.2",
+    version="0.9.3",
     author="Marek Wojciechowski",
     author_email="mrkwjc@gmail.com",
     description="Multithreaded version of numpy.einsum function",


Test patch:
diff --git a/test_einsumt.py b/test_einsumt.py
index b15b104..92d467b 100644
--- a/test_einsumt.py
+++ b/test_einsumt.py
@@ -123,6 +123,12 @@ class TestEinsumt(object):
         A2 = np.einsum('egmipq, egpqrs -> egmirs', u, A)
         k2 = np.einsum('egmirs, egnjrs -> eminj', A2, v)
         assert np.allclose(k1, k2)
+    
+    def test_broadcasting(self):
+        a = np.random.rand(3, 3, 1, 1)
+        b = np.random.rand(3, 3, 6, 100)
+        subs = 'ijrs,ijrs->rs'
+        assert np.allclose(np.einsum(subs, a, b), einsumt(subs, a, b))
 
 
 if __name__ == '__main__':