How to make vmap work with a namedtuple as output of a function?

[shailesh1729]

Following is my example code. It has two functions f and g. Vectorizing f is easy. The output of g is a named tuple. I am unable to figure out how to set the correct value of out_axes for vectorizing g.

import jax
import jax.numpy as jnp

def f (x, y):
    return x+y, x-y

x = jnp.array([1,2,3,4])
y = jnp.array([11,12,13,14])


vf = jax.vmap(f, in_axes=(0, 0), out_axes=(0, 0))

sums, differences = vf(x, y)
print(sums)
print(differences)

from typing import NamedTuple

class SD(NamedTuple):
    u:int = 0
    v:int = 0

def g(x, y):
    return SD(u=x+y, v=x-y)

print(g(2,3))

vg = jax.vmap(g, in_axes=(0, 0), out_axes=(0, 0))

results = vg(x, y)

I get the following error.

ValueError: vmap out_axes specification must be a tree prefix of the corresponding value, got specification (0, 0) for value tree PyTreeDef(namedtuple[<class '__main__.SD'>], [*,*])

I had tried with dataclass too. I faced similar problems. How is the handling of NamedTuple different from a regular tuple in vmap?

[jakevdp]

The out_axes should either be an integer, or have the same structure as the variables. So I believe both out_axes=0 and out_axes=SD(0, 0) will do what you wish.

[shailesh1729]

Thanks. This works perfectly.