How to make pmap return a collection instead of a concatenation

[sokol11]

Hi. I'm running a jit compiled fn with pmap in a distributed fashion with different seeds. The compiled fn returns a python dict. It works great but pmap returns one big dict with concatenated arrays inside, instead of a list of individual dicts, where each dict corresponds to one of the parallel task. The structure of the returned dictionary is somewhat complicated, it contains multiple arrays and other nested dicts with arrays, so splitting it in post-processing is cumbersome and error-prone. Is there a way to make pmap output a list of dicts instead of one big dict with concatenated arrays. Alternatively, is there a way to make it keep arrays from parallel tasks in collections of some sort that I can index into (as opposed to concatenating them)? Thank you!

[jakevdp]

JAX vectorization and parallelization transformations, including vmap, pmap, and shard_map, all work with a struct-of-arrays storage pattern rather than an array-of-structs pattern. This means that if you vmap or pmap over a dict, you get a dict of batched arrays, not a list of dicts. There is no way to make these transformations return an array of structs, but you could take the output and transform it as a post-processing step.

For example, imagine your pmap created this nested dict of arrays with a leading batch dimension of size 4; you could use jax.tree.transpose to convert it to a sequence of dicts:

import jax
import jax.numpy as jnp

result = {'a': jnp.arange(4), 'b': {'c': jnp.ones((4, 2)), 'd': jnp.zeros((4, 7))}}

jax.tree.transpose(jax.tree.structure(result), None, jax.tree.map(list, result))

[{'a': Array(0, dtype=int32),
  'b': {'c': Array([1., 1.], dtype=float32),
   'd': Array([0., 0., 0., 0., 0., 0., 0.], dtype=float32)}},
 {'a': Array(1, dtype=int32),
  'b': {'c': Array([1., 1.], dtype=float32),
   'd': Array([0., 0., 0., 0., 0., 0., 0.], dtype=float32)}},
 {'a': Array(2, dtype=int32),
  'b': {'c': Array([1., 1.], dtype=float32),
   'd': Array([0., 0., 0., 0., 0., 0., 0.], dtype=float32)}},
 {'a': Array(3, dtype=int32),
  'b': {'c': Array([1., 1.], dtype=float32),
   'd': Array([0., 0., 0., 0., 0., 0., 0.], dtype=float32)}}]

That said, you might think about whether this is what you really want: after all, any other JAX-transformed operation you do on this will also expect the data to be in its original struct-of-arrays form, so you may find this array-of-structs version harder to work with.

If instead you want to extract one batch at a time, you can do so using jax.tree.map over the appropriate indexing operation; e.g.

jax.tree.map(lambda x: x[1], result)

{'a': Array(1, dtype=int32),
 'b': {'c': Array([1., 1.], dtype=float32),
  'd': Array([0., 0., 0., 0., 0., 0., 0.], dtype=float32)}}

[sokol11]

Hey! Thanks so much! These one-liners give exactly what I wanted, thanks!