Implementing Nested Loops With Multiple Shape And Size Inputs

[dominicpasquali]

Hello,
I have a JAX function which takes multiple 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:
My inputs look like:

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

input2 = jnp.array([4,5,6,7,8,9])
print(input2.shape) -> (6,)

input3 = jnp.array([[8,9,10]])
print(input3.shape) -> (1,3)

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

out = MyFunction(input1[0], input2, input3)
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 inputs are of the shape

print(input1.shape) -> (512, 100, 2, 16)
print(input2.shape) -> (16,)
print(input3.shape) -> (1,2)

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], input2, input3)))

        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.

Working off of this previous question I tried playing around with the in_axes parameter for vmap:

result = jax.vmap(jax.vmap(jax.vmap(sample_circuit, in_axes = (None, 16, 1))))(input1, input2, input3)
return result

but I keep running into errors of the sort:

ValueError: vmap got inconsistent sizes for array axes to be mapped:
arg 0 has shape (512, 100, 2, 16) and axis 0 is to be mapped
arg 1 has shape (16,) and axis 0 is to be mapped
arg 2 has shape (1, 2) and axis 0 is to be mapped
so
arg 0 has an axis to be mapped of size 512
arg 1 has an axis to be mapped of size 16
arg 2 has an axis to be mapped of size 1

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]

I think you want:

jax.vmap(
    jax.vmap(
        jax.vmap(
            MyFunction,
            (0, None, None)
        ),
        (0, None, None)
   ),
   (0, None, None)
)(inputs, input2, input3)

This is saying that the first 3 axes of inputs should be vmapped over, but input2 and input3 should be broadcast up to the shape of inputs.

[dominicpasquali]

@davisyoshida, thanks for the response! I think there's something I'm not understanding- shouldn't it be:

jax.vmap(
    jax.vmap(
        jax.vmap(
            MyFunction,
            (0, 0, 0, 0)
        ),
        (0, None, None)
   ),
   (0, None, None)
)(inputs, input2, input3)

because according to vmap documentation we need to map the tuple to "which axis to map for each corresponding positional argument" (hence 4 None's vs 3 None's). Also with None "indicating not to map any axis," shouldn't we want each axis for the input1 to map to their correct respective axis?

I apologize for the questions- I'm still trying to fully understand what is actually going on.

[YouJiacheng]

@dominicpasquali there are 3 positional arguments, why use 4 None's?
Both davi and I follow the loop you give here:

 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], input2, input3)))

[dominicpasquali]

@YouJiacheng - you are totally right and I'm completely wrong. I was thinking about the number 4 because I was getting mixed up in that print(input1.shape) -> (512, 100, 2, 16) has 4 dimensions, but vmap is not referring to each individual dimensional shape of an input but rather about each individual input. My apologizes!

[YouJiacheng]

@davisyoshida ’s solution is correct, but I suggest using closure(or partial) to simplify the code.

def foo(inputs, input2, input3)
    @jax.vmap
    @jax.vmap
    @jax.vmap
    def f(inputs_unmapped):
        return MyFunction(inputs_unmapped, input2, input3)
    return f(inputs)

[davisyoshida]

Ah yeah that's nice

[dominicpasquali]

@YouJiacheng - out of curiosity what would a solution with partial look like?

[YouJiacheng]

@dominicpasquali well, it will need MyFunction accept input2 and input3 as keyword arguments.

from functools import partial
def foo(inputs, input2, input3):
    f = partial(MyFunction, input2=input2, input3=input3)
    f = jax.vmap(jax.vmap(jax.vmap(f)))
    return f(inputs)

[davisyoshida]

Or you can partial the vmap itself:

from functools import partial

batch_first = jax.partial(vmap, in_axes=(0, None, None))
result = batch_first(batch_first(batch_first(MyFunction)))(inputs, input2, input3)

I don't think there's any difference in the final outcome so it should all just be down to taste.

[dominicpasquali]

@YouJiacheng and @davisyoshida - if I could mark both of your answers as "the answer," I would. I clearly have a lot to learn and I'm grateful for having both of you to help me out. Thank you for all of your help!