How to specify in_axes for custom pytree node?

[IrishWhiskey]

Let's say I have a function, which takes a custom object as input, that I want to transform with pmap. How should the in_axes argument look like?

If I run

import jax
import jax.numpy as jnp
from functools import partial

@jax.tree_util.register_pytree_node_class
class C:
    def __init__(self, data, static_data):
        self.data = data
        self.static_data = static_data
        
    def tree_flatten(self):
        children = (self.data,)  # arrays / dynamic values
        aux_data = dict(static_data=self.static_data)  # static values
        return (children, aux_data)

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

@partial(jax.pmap, in_axes=(C(0, None)))
def foo(x):
    return jnp.exp(x.data)

c = C(jnp.ones((100, 100)), 1)
foo(c)

I get ValueError: pmap in_axes specification must be a tree prefix of the corresponding value, got specification (<__main__.C object at 0x7f26fedbd0c0>, 0) for value tree PyTreeDef(((CustomNode(C[{'static_data': 1}], [*]),), {})).

Am I doing anything wrong? How can I fix it?

[patrick-kidger]

You want in_axes=(C(0, 1),).

• in_axes should be a tuple of length equal to your number of input arguments.
• the static data for the custom pytree node must be the same.

[davisyoshida]

Thanks for the response, I'll keep that in mind. It's pretty hard to resist getting fancy with them given how useful they are though!

[darsnack]

Apologies for hijacking this thread after many months, but I've been searching for quite some time on this issue, and this answer is the closest discussion I've found.

So, following up on above, the requirement that static leaves are the same means that it isn't possible to specify out_axes without already having an instance of the object? This precludes vmapping class creation methods over seeds as an example.

[jakevdp]

I think that's correct: in order to specify non-trivial out_axes, you need to pass an instance of the class. If you're unable to construct an instance of the class for some reason, you will not be able to specify non-trivial out_axes. That said, I don't totally understand the context of what you describe: I'd suggest opening a new issue with a minimal repro of the kind of thing you have in mind.

[patrick-kidger]

Getting an instance of an object is doable by first calling jax.eval_shape :)

(And then modifying it with e.g. jax.tree_map and equinox.tree_at to obtain your desired out_axes.)

[jakevdp]

Whether eval_shape is a solution depends on the details of the setting the user describes. Let's move this discussion to a new issue if @darsnack would like to dig-in more.