Tensor function optimization with Jax

[MRiabov]

Sup Jax,

I have an interesting usecase - I have a struct that contains a 4-dimensional tensor with millions of elements - in which tensor contains 17 physical properties of elements.

On each of those physical properties, or a set of them, a function is done. For example, gravity(), computeVelocity(), computeStrain(), etc.

In Jax, would it be more efficient to pass references to the entire axis of a tensor and fetch data in the functions, or would it be more efficient to fetch each outside of the functions (in the call)

I use a 4-dimensional vmap for computation, like `vmap(vmap(vmap(vmap(fun))))(a,b,c)

To clarify:

vmap(vmap(vmap(vmap(mega_function)))(tensor)

def mega_function(tensor)
   func1(tensor.at[0,1,2]
   func2(tensor.at[3,4,5,6,7,8])
   func3(tensor.at[9,10,11,12,13,14]
   ....

or

vmap(vmap(vmap(vmap(func1)))(tensor[0],tensor[1],tensor[2])
vmap(vmap(vmap(vmap(func2)))(tensor[3],tensor[4],tensor[5],tensor[6],tensor[7],tensor[8])
...

on them?

How is this actually managed in jax? would there be any difference?
Thanks.

[jakevdp]

You'll probably find that the second approach will be more efficient, because it avoids mapped indexing operations over the arrays.