Slow code with Jax

[raj-brown]

@mattjj Hi Matt-
I wanted to implement a prototype using forward mode gradient, using jvp make the code very slow.
I am very new to jax and will appreciate if you can help me out. Below is the MWE . Thank you very much! - Raj

[YouJiacheng]

The code is slow because of compiling.
jax.jit will completely unroll python for-loop, and you write a 10k step for-loop inside update.
That will make JAX compile extremely slow.

[raj-brown]

@YouJiacheng Thank you! @mattjj @YouJiacheng could you please help on how to use vmap for the following

   v0=  random.normal(subkey, s_array[0])
    v1 = random.normal(subkey, s_array[1])
    v2 = random.normal(subkey, s_array[2])
    v3 = random.normal(subkey, s_array[3])
    v4 = random.normal(subkey, s_array[4])
    v5 = random.normal(subkey, s_array[5])

Thank you very much!

[jakevdp]

Hi - there is currently no way to use vmap for such an operation, because in general it would create a ragged array (i.e. an array in which each row has a different length). JAX does not support ragged arrays, and so it cannot support the operation you have in mind via vmap.

Your best bet is probably something like this:

v = [random.normal(subkey, s) for s in s_array]

though you should probably avoid using the same subkey for each row.