It's not true that einsum converts to a for loop, and it is sometimes faster than other numpy built-in functions.
https://stackoverflow.com/questions/18365073/why-is-numpys-e...
Did I misunderstood your comment?
https://stackoverflow.com/questions/18365073/why-is-numpys-e...
Did I misunderstood your comment?
The optimizer clearly tries to improve the performance, but in many cases, it doesn't seem to change anything. Let's simply multiply some matrices:
x, y = np.random.rand(200, 200, 200), np.random.rand(200, 200, 200)
%timeit x@y
40.3 ms ± 2.52 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit np.einsum('bik,bkj->bij',x,y)
1.53 s ± 21.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit np.einsum('bik,bkj->bij',x,y, optimize=True)
1.54 s ± 10.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit (x_jax @ y_jax).block_until_ready()
579 µs ± 4.54 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit jnp.einsum('bik,bkj->bij',x_jax,y_jax, optimize=True).block_until_ready()
658 µs ± 1.38 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit jnp.einsum('bik,bkj->bij',x_jax,y_jax).block_until_ready()
660 µs ± 2.82 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)`optimize=True` is generally best when there are more than two tensors in the expression.