How to write a partial vmap with a user PyTree

[jecampagne]

Hello, let me write some snippets to introduce my problem on applying vmapon a user PyTree.
So, first let me practice with a dictionary:

from jax import vmap, grad
def f(p,x):
    return p["a"]**2 + p["b"]*x
grad(f)({"a":3.,"b":1.},10)
vmap(f, in_axes=({"a": None, "b": 0},None))({"a":1.,"b":jnp.array([1.,2.,3.])},10)

this gives respectively

{'a': DeviceArray(6., dtype=float64), 'b': DeviceArray(10., dtype=float64)}
[11. 21. 31.]

So far so good, now let me define my PyTree

from jax.tree_util import register_pytree_node_class

@register_pytree_node_class
class Params:
    def __init__(self, a, b):
        self._a = a
        self._b = b

    def __repr__(self):
        return f"Params(x={self._a}, y={self._b})"

    @property
    def a(self):
        return self._a
    
    @property
    def b(self):
        return self._b

    def tree_flatten(self):
        children = (self._a, self._b)
        aux_data = None
        return (children, aux_data)

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

and play with it the same game as with the dictionnary

def g(p,x):
    return (p.a)**2 + (p.b)*x

my_params = Params(3.,1.)
print(grad(g)(my_params,10))

gives

Params(x=6.0, y=10.0)

ok, but how to write the equivalent of in_axes=({"a": None, "b": 0},None)?

Thanks.

[YouJiacheng]

Just treat your pytree as a container.

import jax
import jax.numpy as jnp

from jax.tree_util import register_pytree_node_class

@register_pytree_node_class
class Params:
    def __init__(self, a, b):
        self._a = a
        self._b = b

    def __repr__(self):
        return f"Params(x={self._a}, y={self._b})"

    @property
    def a(self):
        return self._a
    
    @property
    def b(self):
        return self._b

    def tree_flatten(self):
        children = (self._a, self._b)
        aux_data = None
        return (children, aux_data)

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


def g(p,x):
    return (p.a)**2 + (p.b)*x

my_params = Params(3.,1.)
print(jax.vmap(g, in_axes=(Params(None, 0), None))(Params(3., jnp.array([1.,2.,3.])), 10))
# [19. 29. 39.]

[jecampagne]

Thanks!