A basic introduction to NumPy's einsum

Post by savant_penguin » Sun, Apr 10, 2022, 12:06 AM UTC

thrownawayalre wrote: post #6

cschmid wrote: post #2

I'd really like to use einsum more often, because it allows me to code my expressions the same way I derive them on pen and paper. Unfortunately, as mentioned in the article, it's slow, because it converts your formula to a for loop. So usually, I rewrite my formulas into messy combinations of broadcasts, transposes and array multiplications. Is there a package or an algorithm that does this conversion automatically?…

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?

Post by madisonmay » Sun, Apr 10, 2022, 12:12 AM UTC

cschmid wrote: post #2

I'd really like to use einsum more often, because it allows me to code my expressions the same way I derive them on pen and paper. Unfortunately, as mentioned in the article, it's slow, because it converts your formula to a for loop. So usually, I rewrite my formulas into messy combinations of broadcasts, transposes and array multiplications. Is there a package or an algorithm that does this conversion automatically?…

Post by cschmid » Sun, Apr 10, 2022, 12:20 AM UTC

cschmid wrote: post #12

thrownawayalre wrote: post #6

Earlier quoted context omitted.

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?

If I understand https://github.com/numpy/numpy/blob/v1.22.0/numpy/core/einsu... and https://github.com/numpy/numpy/blob/v1.22.0/numpy/core/src/m... correctly, using einsum without the optimize flag seems to use a for loop in C to do the multiplication. 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.ra…

  %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)

Post by upbeat_general » Sun, Apr 10, 2022, 1:09 AM UTC

enriquto wrote: post #7

cschmid wrote: post #2

I'd really like to use einsum more often, because it allows me to code my expressions the same way I derive them on pen and paper. Unfortunately, as mentioned in the article, it's slow, because it converts your formula to a for loop. So usually, I rewrite my formulas into messy combinations of broadcasts, transposes and array multiplications. Is there a package or an algorithm that does this conversion automatically?…

Just use C or C++ or fortran or julia or even lua. The issue of "slow loops" is entirely self-inflicted by some languages. It's getting ridiculous to still worry about this shit in 2022. Simple loops can and should be just as fast as vectorized programs. When they are slower, it is 100% due to deliberate decisions by the language dessigners

Post by catgary » Sun, Apr 10, 2022, 1:28 AM UTC

cschmid wrote: post #2

I'd really like to use einsum more often, because it allows me to code my expressions the same way I derive them on pen and paper. Unfortunately, as mentioned in the article, it's slow, because it converts your formula to a for loop. So usually, I rewrite my formulas into messy combinations of broadcasts, transposes and array multiplications. Is there a package or an algorithm that does this conversion automatically?…

  @tullio C[i,j] := A[i,k] * B[k,j]