Efficiently implementing nested loops with single array JAX function

[dominicpasquali]

Hello,
I have a JAX function which only takes single JAX NumPy (JNP) arrays (of the shape (number,) ), it outputs single JNP arrays, and I'm trying to figure out how to incorporate this function into my __call__ function for my ML module.

For example this function will behave like the following:
If my input looks like:

input = jnp.array([[0,1,2,3]])
print(input.shape) -> (1,4)
print(input[0]) -> [0,1,2,3]
print(input[0].shape) -> (4,)

then my function will behave like (note that the below is just a dummy output to give an example)

out = MyFunction(input[0])
print(out) -> [4,5,6,7]
print(out.shape) -> (4,)

The problem arises when I start handling multidimensional arrays for my input. For example my input is of the shape

print(input.shape) -> (512, 100, 2, 16)

I need to make sure of 3 things: I iterate over each input, I store each output, and I reformat my output to be properly formatted for other ML modules. To do this I run:

@nn.compact
    def __call__(self, inputs):

        Zero_Array = jnp.zeros(shape=inputs.shape)

        for L in range(inputs.shape[0]):
            for M in range(inputs.shape[1]):
                for N in range(inputs.shape[2]):
                    Zero_Array = Zero_Array.at[L, M, N].set(jnp.array(MyFunction(inputs[L][M][N])))

        return Zero_Array

As you can imagine my code runs very slow and very quickly runs out of RAM. I'm struggling to implement more efficient loops with vmap, lax.fori_loop, or lax.scan but I'm not quite sure how to do it since I'm overwriting Zero_Array each time but I want to preserve inputs. Also I considered using the loops module, but it's being deprecated later this year.

Any ideas on how to make the above code more efficient and/or rewrite it using vmap, lax.fori_loop, and/or lax.scan?

Thanks!

(if I get a solution on my own I'll update this thread accordingly)

[davisyoshida]

There may be some way to use a single vmap call, but you can definitely do:

result = jax.vmap(jax.vmap(jax.vmap(MyFunction)))(inputs)

By default, each call to vmap is adding another batch dimension to the left of the shape of the inputs/outputs.

[dominicpasquali]

Hi @davisyoshida - thanks for your answer! I just realized that I messed up on the question as I need to pass multiple different size inputs to my function. Since what my question should have been is different enough from the current question, (and I'm still stuck on it) I'll open up another question. I'll leave this question as is for others who find this response helpful.

[YouJiacheng]

For multiple different size inputs, you can use in_axes to specify which input is mapped in this vmap. e.g.
vmap(vmap(f), in_axes=(0, None)) will map 2 axes of the first argument, and 1 axis of the second argument.

[dominicpasquali]

@YouJiacheng I tried doing that with no luck, though I'm sure it's something as simple as that (I think I'm thinking about it the wrong way, and I think I'm placing the in_axes parameter at the wrong location)

[dominicpasquali]

my new question can be found here: #10928

[YouJiacheng]

Two simple solution: Reshape input and use single vmap or just use nested vmap.

[dominicpasquali]

Nested vmap is the way to go! Though As I just commented I messed up on the question.