vmap over a list of class

[jing-alice]

vmap does not currently fit into a list of classes and will automatically expand all the attributes of the class, causing an error in np.shape(x) in vmap.
for example:

class Comp:
   def __init__():
       pass
   def __call__():
      pass

self.comp_list =[Comp(), Comp(), Comp()]

dvar_list = vmap(lambda var, comp: comp(var)))(jnp.array([0., 0., 0.]), self.comp_list)

Also, I have tried the register_pytree_node in jax.tree_util, but it doesn't work.

[jakevdp]

Hi - thanks for the question!

Your list of classes is what's known as an array-of-structs pattern. On the other hand, vmap is built to work with a struct-of-arrays pattern. So rather than defining a structure like this:

data_array_of_structs = [Comp(1), Comp(2), Comp(3)]

you should instead define a structure like this:

data_struct_of_arrays = Comp(jnp.array([1, 2, 3]))

Then, as long as Comp is registered as a pytree, you can do this:

jax.vmap(lambda comp: comp())(data_struct_of_arrays)

Hope that helps!