How to vectorize functions of a custom pytree?

[langmore]

I'm trying to use a custom pytree in a differential equation. I get errors when I vectorize. What am I doing wrong?

import jax
from jax.tree_util import register_pytree_node_class

@register_pytree_node_class
class State:
  def __init__(self, t, x):
    self.t = t   # Scalar
    self.x = x   # Vector

  def tree_flatten(self):
    children = (self.t, self.x)
    aux_data = None
    return (children, aux_data)

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

def dsdt(state):
  return State(t=1, x=state.t * state.x)

axes = State(t=None, x=0)
vec_dsdt = jax.vmap(dsdt, in_axes=axes, out_axes=axes)
vec_dsdt(State(t=1, x=jnp.array([1, 2.])))

The traceback includes

ValueError                                Traceback (most recent call last)
[<ipython-input-14-2ba90acfbae3>](https://colab.corp.google.com/drive/1BRozCo_GrxmmgnahDkHl6koUZ53_z5SZ#) in <module>()
     22 axes = State(t=None, x=0)
     23 vec_dsdt = jax.vmap(dsdt, in_axes=axes, out_axes=axes)
---> 24 vec_dsdt(State(t=1, x=jnp.array([1, 2.])))

ValueError: vmap in_axes specification must be a tree prefix of the corresponding value, got specification <__main__.State object at 0x7fc08c2c5310> for value tree PyTreeDef((CustomNode(<class '__main__.State'>[None], [*, *]),)).

[jakevdp]

You should call it like this:

vec_dsdt = jax.vmap(dsdt, in_axes=(axes,), out_axes=axes)

in_axes needs to be a tuple (because in general *args for a function is a tuple, even if there is a single arg). This is perhaps slightly unclear because for convenience in the common case, JAX allows passing a single integer in place of a length-1 tuple.