Implementing vmap over list of PyTrees

[TeodorPoncu]

Hey there! I'm currently trying to build a didactic example for a class I'm holding on Neural Rendering.

I want to focus less on numerical / implementation optimizations in order to make the graphics concepts more easy to follow. At the same time I want to illustrate some of the nice things about JAX (jit, grad, vmap, pmap, etc.). Therefore, I went for an OOP like approach - loosely inspired from this google-research repo.

I have searched for similar issues(#5322) and am aware that a Struct-of-Arrays is preferred in terms of efficiency. I am just curious if there is any sensible way of implementing this toy example without making use pack-unpack tricks (I feel that manually flattening the tree and unflattering the tree inside some specific ops would defeat the purpose of the example and I should just directly resort to SoA variants)

from typing import Union
import jax.numpy as jnp
import jax

Numeric = Union[float, int, complex, jnp.ndarray]

@jax.tree_util.register_pytree_node_class
class Point3:
  def __init__(self, x, y, z, dtype=jnp.float32) -> None:
    self.arr = jnp.asarray([x, y, z], dtype=dtype)
    
  def x(self) -> Numeric:
    return self.arr[0]
  
  def y(self) -> Numeric:
    return self.arr[1]
  
  def z(self) -> Numeric:
    return self.arr[2]
  
  def __add__(self, other: "Point3") -> "Point3":            
    return Point3(*(self.arr + other.arr))
  
  def tree_flatten(self):
    children = self.arr
    aux_data = None
    return (children, aux_data)
  
  @classmethod
  def tree_unflatten(cls, aux_data, children):
    del aux_data
    return cls(*children)
  
  def __repr__(self) -> str:
    return f"Point3({self.x()}, {self.y()}, {self.z()})"

My goal for one exercise would be to vmap on top of a list of such primitives where I make use of the overloaded operators.

if __name__ == "__main__":
  p1 = Point3(1, 2, 0)
  p2 = Point3(3, 4, 1)

  line_1 = [p1, p2]

  line_2 = jax.vmap(lambda x: x + x)(line_1)

However, when trying to vmap that function I get the following error:

 out = np.array(object, dtype=dtype, ndmin=ndmin, copy=False)
jax._src.traceback_util.UnfilteredStackTrace: TypeError: float() argument must be a string or a real number, not 'object'

When inspecting what gets passed in from the tree_unflatten operation into the constructor

@jax.tree_util.register_pytree_node_class
class Point3:
  def __init__(self, x, y, z, dtype=jnp.float32) -> None:
    print(x, y, z)
    self.arr = jnp.asarray([x, y, z], dtype=dtype)
...
>>> (<object object at 0x7f015c065390>, <object object at 0x7f015c065390>, <object object at 0x7f015c065390>)

The full stack-trace is the following:

Traceback (most recent call last):
  File "test.py", line 44, in <module>
    line_2 = jax.vmap(lambda x: x + x)(line_1)
  File "env/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "env/site-packages/jax/_src/api.py", line 1677, in vmap_f
    in_axes_flat = flatten_axes("vmap in_axes", in_tree, (in_axes, 0), kws=True)
  File "env/site-packages/jax/_src/api_util.py", line 393, in flatten_axes
    dummy = tree_unflatten(treedef, [object()] * treedef.num_leaves)
  File "env/site-packages/jax/_src/tree_util.py", line 75, in tree_unflatten
    return treedef.unflatten(leaves)
  File "test.py", line 33, in tree_unflatten
    return cls(*children)
  File "test.py", line 10, in __init__
    self.arr = jnp.asarray([x, y, z], dtype=dtype)
  File "env/site-packages/jax/_src/numpy/lax_numpy.py", line 2028, in asarray
    return array(a, dtype=dtype, copy=False, order=order)  # type: ignore
  File "env/site-packages/jax/_src/numpy/lax_numpy.py", line 1990, in array
    out = np.array(object, dtype=dtype, ndmin=ndmin, copy=False)
jax._src.traceback_util.UnfilteredStackTrace: TypeError: float() argument must be a string or a real number, not 'object'

Sorry if I broke any discussions formatting rules!

[jakevdp]

vmap can only be used to map over arrays, not over lists. And arrays in JAX can only contain numerical values, not structures like pytrees. So in your case, if you want to use vmap it would look something like this:

points = Point3(jnp.array([1, 2]), jnp.array([2, 4]), jnp.array([0, 1]))
jax.vmap(lambda x: x + x)(points)

As you mentioned, this is a struct-of-arrays pattern, whereas your example was an array-of-structs pattern. If you're unable to switch to a struct-of-arrays pattern, you won't be able to use vmap. But you could use other approaches, for example Python's built-in map function works well with lists.

Regarding the error you're seeing: it's not directly related to the struct-of-arrays issue, but rather to the fact that to be used within JAX transforms, PyTrees must be flexible regarding input validation. This issue is discussed at Custom PyTrees and Initialization.