vmap against list of pytree containers

[bhchiang]

Let's say I want to vmap a function against a list of pytree objects. For example, here I have a namedtuple as my pytree.

import jax.numpy as jnp
from jax import vmap
from collections import namedtuple

def trace(sphere):
    # trace ray against sphere
    pass

Sphere = namedtuple("Sphere", ["center", "radius"])
surfaces = jnp.array([
    Sphere(jnp.zeros(3), 1),
    Sphere(jnp.array([0, 2, 0]), 10)
])
# ValueError: All input arrays must have the same shape.

results = vmap(trace)(surfaces)

This entire snippet is running inside a vmap, and ideally I want to do everything in parallel - but I can't create the list of surfaces as a jnp.ndarray.

Other options I was thinking about:

• lax.scan / lax.fori_loop - but I don't think I can iterate over a plain Python list.
• experimental.loops and index into each sphere
• Manually pack the pytree into a flat jnp.ndarray via jnp.hstack(tree_flatten(sphere)), and then manually reconstruct the tree via tree_unflatten inside trace
• tree_map?

Any suggestions appreciated, thanks!

[shoyer]

vmap supports pytrees when written in the form of a "struct of arrays" rather than an "array of structs". This is because JAX/XLA only has support for numeric array dtypes, so the "array of struct" form can't be operated on efficiently in JAX.

In other words, you need to instead move the array dimension inside the containers, e.g.,

centers = jnp.stack([jnp.zeros(3), jnp.array([0, 2, 0])])
radiuses = jnp.stack([1, 10])
surfaces = Sphere(centers, radiuses)

Then vmap(trace)(surfaces) should work just fine, returns a namedtuple of arrays.

[bhchiang]

Thanks for clearing this up - I'll try writing a custom stacked_vmap function that does this automatically for a list of data classes (registered as pytrees).

[vaishnkv]

@bhchiang ,
Can you pls share the code or pseudo-code for the stacked_vmap , if you have implemented the same?

[bantin]

@shoyer I have the same issue, and I'm wondering about best practices for a scenario like this. In my case, I have a list of NamedTuples. I have functions that operate on these NamedTuples and return devicearrays. For readability, I'd like to avoid re-writing those functions.

For instance:

Point = namedtuple("point", ["x", "y"])

@jax.jit
def f(pt):
    return pt.x**2 + pt.y**2

vmap(f)((Point(1,2), Point(2,3)))

won't work out of the box, as you clarified above. My question is whether there's a way of getting this to work without modifying the function f in my example? Can I do something like a tree_transpose here? In the scenario described above, it looked like we would have had to rewrite the trace function, which is something I need to avoid.

[bhchiang]

What you can do is stack all of the Points into a single NamedTuple using tree_multimap.

Example here: https://github.com/bryanhpchiang/rt/blob/master/utils/transforms.py#L5.

[jecampagne]

Hello, @bhchiang
How do you use your nice tool?
Let say that I have a list of Pytree objects

from jax.tree_util import register_pytree_node_class

@register_pytree_node_class
class A():
    def __init__(self,val, f):
        if not (type(val) is object or val is None or isinstance(val, A)):
            val = jnp.asarray(val)
        self.val = val

        self.f = f

    def __repr__(self):
      return "A(val={})".format(self.val)
    
    def tree_flatten(self):
      children = (self.val,self.f)
      aux_data = None
      return (children, aux_data)

    @classmethod
    def tree_unflatten(cls, aux_data, children):
      return cls(*children)

build with the following facztory ( notice the Partial)

from jax.tree_util import Partial

class A_factory():
    done = False
    _ws = {}
    @classmethod
    def make(cls, n_data=5_000):
        if not A_factory.done:
            A_factory._ws = {'data': [A(val, Partial(K)) for val in np.random.uniform(0., 10, size=n_data)]}   #### notice the Partial
            A_factory.done = True
        return A_factory._ws
list_of_A =A_factory.make()['data']

and a function

def func(a_obj, x):
  return a_obj.val * a_obj.f(x)

For the tile beeing I am using

def f(p):
    list_of_A = A_factory.make()['data']
    data = jnp.array([func(a,p) for a in list_of_A])
    return data

but I wander if I can use your tool to be able to use a vmap(func)(list_of_A) ? Notice that the jnp.array in this context is quite slow but I need it for autodiff wrt 'p' argument.

[jecampagne]

I seems that I can manage to perform

@jit
def test(p):
  data = jax.vmap(lambda x: func(x,p))(pytrees_stack(list_of_A))
  return data

Is it correct? (I do not manage to use your pytrees_vmapcertainly because I miss something.
Notice that your method may be solution of my problem #11324