Trying to pmap a list of PyTree

[ThatBitBoss]

Here is my code

When I go over an array of integers - it seems to work

But the list of PyTree seems to throw an error.
Ultimately I want to pmap to a execusteGoal_Hunger and I am obviously missing or messing up something.

[jakevdp]

Hi - thanks for the question! pmap operates over arrays, not over lists. When collections like lists, tuples, or dicts are passed to JAX functions like pmap, vmap, map, scan, etc. the semantics are that these function iterate over all arrays in the collection.

So to do what you wish here, instead of passing a list of pytrees, you should pass a pytree of arrays. Here's a simple example of what I mean, using vmap for simplicity:

import jax

def f(x):
  return x['A'] + x['B']

list_of_trees = [
    {'A': 1, 'B': 2},
    {'A': 5, 'B': 10}
]
jax.vmap(f)(list_of_trees)  # wrong!
# ValueError: vmap was requested to map its argument along axis 0, which implies
#             that its rank should be at least 1, but is only 0 (its shape is ())

tree_of_arrays = {'A': jnp.array([1, 5]), 'B': jnp.array([2, 10])}
result = jax.vmap(f)(tree_of_arrays)  # correct!
print(result)
# [ 3 15]

Your code uses a list_of_trees pattern, instead you should use a tree_of_arrays pattern if you want the function to be compatible with pmap or vmap.

Hope that helps!