jnp.einsum does not catch shape mismatch when optimize=True

There seems to be a bug in jnp.einsum. The function does not catch a shape mismatch when the optimize flag is set to True.

import jax.numpy as jnp

A = jnp.ones(2)[:,None]
x = jnp.ones(2)
try:
    A@x
    print('Passed?')
except:
    print('Failed as expected')

try:
    jnp.einsum('ij,j->i', A, x, optimize=False)
    print('Passed?')
except:
    print('Failed as expected')

try:
    out = jnp.einsum('ij,j->i', A, x, optimize=True)
    print('Passed?')
    print(out)
except:
    print('Failed as expected')
    
# Failed as expected
# Failed as expected
# Passed?
# [2. 2.]

[mattjj]

Thanks for catching this @EddieCunningham , and wow nice digging @IgorWilbert ! So should we consider this a bug in opt_einsum?

Maybe we should do our own shape checking.

[shoyer]

I think both NumPy and JAX should be strict about requiring exactly matching sizes. The entire point of einsum is to avoid implicit broadcasting…