The einsum-experiment from sorenlassen

Broadcast problem

I was testing the examples at https://www.tensorflow.org/api_docs/python/tf/einsum and noticed an interesting failure case (search for "Einsum will broadcast over axes covered by the ellipsis" on that page):

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

The example right above it works fine:

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.

sorenlassen / einsum-experiment Goto Github PK

einsum-experiment's People

Contributors

Stargazers

Watchers

Forkers

einsum-experiment's Issues

Broadcast problem

Recommend Projects

React

Vue.js

Typescript

TensorFlow

Django

Laravel

D3

Recommend Topics

javascript

web

server

Machine learning

Visualization

Game

Recommend Org

Facebook

Microsoft

Google

Alibaba

D3

Tencent

Jobs