import tensorflow as tf
import einsum
s = tf.random.normal(shape=[11, 1, 5, 3])
t = tf.random.normal(shape=[1, 7, 3, 2])
e = tf.einsum('...ij,...jk->...ik', s, t)
print(e.shape)
e2 = einsum.einsum('...ij,...jk->...ik', s.numpy(), t.numpy())
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
File "/Users/shawn/ml/einsum/src/einsum/lib.py", line 228, in einsum
return einsum_execute(spec, list(tensors))
File "/Users/shawn/ml/einsum/src/einsum/lib.py", line 223, in einsum_execute
return einsum_tensor_frompos(fn, spec.output.shape)
File "/Users/shawn/ml/einsum/src/einsum/lib.py", line 31, in einsum_tensor_frompos
return einsum_tensor(recurse([], shape))
File "/Users/shawn/ml/einsum/src/einsum/lib.py", line 29, in recurse
return [ recurse(pos + [i], remainder) for i in range(shape[0]) ]
File "/Users/shawn/ml/einsum/src/einsum/lib.py", line 29, in <listcomp>
return [ recurse(pos + [i], remainder) for i in range(shape[0]) ]
File "/Users/shawn/ml/einsum/src/einsum/lib.py", line 29, in recurse
return [ recurse(pos + [i], remainder) for i in range(shape[0]) ]
File "/Users/shawn/ml/einsum/src/einsum/lib.py", line 29, in <listcomp>
return [ recurse(pos + [i], remainder) for i in range(shape[0]) ]
File "/Users/shawn/ml/einsum/src/einsum/lib.py", line 29, in recurse
return [ recurse(pos + [i], remainder) for i in range(shape[0]) ]
File "/Users/shawn/ml/einsum/src/einsum/lib.py", line 29, in <listcomp>
return [ recurse(pos + [i], remainder) for i in range(shape[0]) ]
File "/Users/shawn/ml/einsum/src/einsum/lib.py", line 29, in recurse
return [ recurse(pos + [i], remainder) for i in range(shape[0]) ]
File "/Users/shawn/ml/einsum/src/einsum/lib.py", line 29, in <listcomp>
return [ recurse(pos + [i], remainder) for i in range(shape[0]) ]
File "/Users/shawn/ml/einsum/src/einsum/lib.py", line 26, in recurse
return fn(*pos)
File "/Users/shawn/ml/einsum/src/einsum/lib.py", line 221, in fn
return recurse(in_only_idxs)
File "/Users/shawn/ml/einsum/src/einsum/lib.py", line 218, in recurse
acc += recurse(tail)
File "/Users/shawn/ml/einsum/src/einsum/lib.py", line 210, in recurse
prod *= tensor.item(*pos)
IndexError: index 1 is out of bounds for axis 1 with size 1
import tensorflow as tf
import einsum
import numpy as np
s = tf.random.normal(shape=[11, 7, 5, 3])
t = tf.random.normal(shape=[11, 7, 3, 2])
e = tf.einsum('...ij,...jk->...ik', s, t)
print(e.shape)
e2 = einsum.einsum('...ij,...jk->...ik', s.numpy(), t.numpy())
np.allclose(e, e2) # True
So the problem seems related to broadcasting axes covered by ellipses.