how to use vmap without collating the output

[aldopareja]

Hi, I'm sorry if this has been asked before, couldn't find an answer.

I have a function that creates an object in the return, I want vmap to return an array of dicts and not a dict of arrays. This code shows what I want and what I got. Any suggestions?

from jax import vmap
from jax.random import uniform, split, PRNGKey

def sample_dict(key:PRNGKey):
  key, subkey = split(key)
  x = uniform(key, (3,))
  y = uniform(subkey, (3,))
  return dict(x=x, y=y)

v_sample_dict = vmap(sample_dict)

out = v_sample_dict(split(PRNGKey(0), 3))
print("what I got")
print(out)

print("what I want")
print([dict(x=x,y=y) for x,y in zip(out['x'],out['y'])])

what I got
{'x': DeviceArray([[0.252136  , 0.2555536 , 0.7349534 ],
             [0.30134034, 0.5773866 , 0.30234015],
             [0.35577416, 0.7246367 , 0.6810206 ]], dtype=float32), 'y': DeviceArray([[0.8332895 , 0.2867911 , 0.64770496],
             [0.027583  , 0.38160014, 0.26059723],
             [0.45792747, 0.3699044 , 0.8316119 ]], dtype=float32)}
what I want
[{'x': DeviceArray([0.252136 , 0.2555536, 0.7349534], dtype=float32), 'y': DeviceArray([0.8332895 , 0.2867911 , 0.64770496], dtype=float32)}, 
{'x': DeviceArray([0.30134034, 0.5773866 , 0.30234015], dtype=float32), 'y': DeviceArray([0.027583  , 0.38160014, 0.26059723], dtype=float32)}, 
{'x': DeviceArray([0.35577416, 0.7246367 , 0.6810206 ], dtype=float32), 'y': DeviceArray([0.45792747, 0.3699044 , 0.8316119 ], dtype=float32)}]

[mattjj]

Great question!

There's no way in JAX to represent an array of dicts; the only element types which are allowed are numeric numpy dtypes like int32, float32, etc.

Another way to think about it is that you have to work in structure-of-arrays form, and you can't have arrays-of-structures.

We've batted around some ideas to let you write arrays-of-structures directly, but we haven't made it a priority. I'd be interested to hear about your use case, and whether writing structures-of-arrays is really painful in it.

[mattjj]

Wow, given that composition of tools, you must be working on something pretty cool!

Well, vmap only works on JAX-based functions. But maybe you have a non-JAX renderer called in the middle of some function which is otherwise implemented in JAX, and you want to batch that function? Something like:

def f(x):
  y = some_jaxy_stuff(x)
  z = non_jax_renderer(y)
  w = more_jaxy_stuff(z)
  return w

where you'd like to be able to use vmap(f), even if it meant the batching of the renderer had to be a loop, or something?

[aldopareja]

yes! even if the batching of the rendering would have to be a loop or something. In general I'm creating a generative model in numpyro that uses a physics engine as a "world model". Then I can use this generative model along with NNs to infer the state of a particular dynamic scene and reason about many possible properties (e.g. navigation, finding objects, predicting collisions, etc).

So my pipeline at the moment would be:

def episode():
    for t in range(num_steps):
        scene_description = numpyro_sample_object_6d_poses()
        observation = non_jax_render(scene_description)
        action = numpyro_NN_sample_action(observation, scene_description)
        scene_decription = brax.step(scene_description, action)

I tried to use the renderer inside the function but, since scene_description is a traced array, I can't use non_jax_renderer because it would throw an error when trying to cast it into numpy. What would be the best way around this?

[mattjj]

@froystig this is a perfect use for custom_batching, and even a convenience wrapper around it to batch via loops.

[mattjj]

Also this is issue/discussion #12345... epic.

[froystig]

Nice. Custom batching tracking issue is #9073.

The loop-based convenience wrapper is already in (via #9275), but custom_vmap doesn't support reverse-mode AD yet.