unbound axis name error using xmap with odeint and SerialLoop

[DanPuzzuoli]

Hi, I'm trying to use xmap to parallelize the mapping of a function over multiple CPU cores and am running into an error.

Below is a small reproduction. I have:

• A function f, whose input and return types are both a single float
• An array input_array of floats of length 16 I want to map f over
• A device with 8 CPU cores

I want to parallelize the mapping of f over input_array using the 8 cores. In this case 2 inputs to f will be evaluated on each core, however I want these to be evaluated using a serial loop. My code attempting this, with a simple version of f, is below:

import os
os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8' # Use 8 CPU devices

import numpy as np
import jax
import jax.numpy as jnp
from jax.sharding import Mesh
from jax.experimental.maps import xmap, SerialLoop
from jax.experimental.ode import odeint

def rhs(y, t):
    return t**2

def f(a):
    return odeint(rhs, y0=jnp.array(a), t=jnp.array([0., 1.]))[-1]

input_array = np.random.rand(16)
devices = np.array(jax.devices())

with Mesh(devices, ("x",)):
    out_array = xmap(
        f,
        in_axes=["map_axis", ...],
        out_axes=["map_axis", ...],
        axis_resources={"map_axis": ("x", SerialLoop(2))}
    )(input_array)

out_array

this yields the error

NameError: unbound axis name: <UniqueResource None 0>. The following axis names (e.g. defined by pmap) are available to collective operations: ['x']

What's confusing me is that if I change the definition of f, e.g. to:

def f(a):
    return a**2

then the code seems to work. Also, if I remove the SerialLoop(2) specification, the code also works. However, for the full version of my function, I definitely want to do a serial loop as it doesn't vectorize well.

I'm not sure here if I'm not using SerialLoop properly, or if there is something about odeint that's not combining well with xmap.

I'm using jax==0.4.8 and jaxlib==0.4.7.

[cgarciae]

This looks like a bug. lax.scan is creating an unnamed resource and xmap doesn't like that.

[itk22]

Hi @cgarciae, do you know if there has been any follow-up on this? I was very excited to find out about the Serial Loop utility but I am encountering the same error with scan and wondering if this is resolved/ will be resolved in near future. Thanks.