Is it possible to look at numpy representation of concrete einops function?

[marinegor]

Hi, first discussion here -- wanted to say big thanks for the package, it's amazing!!

Now, question: from what I understood while reading the docs & code, einops comes up with a pure-numpy/pytorch/jax set of functions for each of its operations. My questions is, is it possible to check which exactly function did it come up with?

I am looking to reduce the dependency number in a library to a bare minimum, and there are 2 places where I use einops -- I wonder if I could just plug a pure-numpy implementation of the function. I know that einops is dependency-less and in principle having it as dependency isn't a big deal, but still.

[arogozhnikov]

Hi @marinegor,
if you want a full-reproduction of einops, it's simpler to just use einops itself, but if you're interested in processing static shapes, I can suggest two options:

1. track this function in debugger (this is where actual calls to numpy/backend happen)

   einops/einops/einops.py

   Line 229 in 5655ce1

   def _apply_recipe(

2. use make_jaxpr from jax to see actually used operations (you'll need to manually convert this to a function), see e.g. discussion in this issue: Support JAX's distributed arrays #252

[marinegor]

Thanks, make_jaxpr does the trick for my case.