grad + vmap + odeint AssertionError

[DanPuzzuoli]

Hi,

I've seen a bunch of discussions and issues surrounding this so I apologize if i'm re-raising something that has already been addressed elsewhere. I don't understand the internals of JAX enough to understand if this is some version of an issue that's already raised, though a lot of what I'm seeing is in already closed issues so I assume are solved and hence this is different.

When I run the following code on the latest release versions of jax/jaxlib (a self-contained version of my actual code):

from jax.experimental.ode import odeint
from jax import jit, value_and_grad, vmap
import jax.numpy as jnp
from jax.config import config
config.update("jax_enable_x64", True)
config.update('jax_platform_name', 'cpu')


T = 1.

X = -1j * jnp.array([[0., 1.], [1., 0.]], dtype=complex)
Y = -1j * jnp.array([[0., -1j], [1j, 0.]], dtype=complex)

def err_obj(a, b_vals):
    
    def err(b):
        def rhs(y, t):
            return (b * X + a * (t**2) * Y) @ y
        
        res = odeint(rhs, y0=jnp.eye(2, dtype=complex), t=jnp.array([0, T], dtype=complex), rtol=1e-6, atol=1e-6)
        
        return jnp.abs((X * res[-1]).sum())**2 / 4


    all_err = vmap(err)(b_vals)
    return all_err.sum()

b_vals = jnp.array([1., 2., 3., 4., 5.])
jit(value_and_grad(lambda a: err_obj(a, b_vals)))(1.)

I get the error:

AssertionError: length mismatch: [1, 4]

I've definitely reverse-mode differentiated this code in the past with success, though it was some time ago so would have been on a much older version of jax/jaxlib. Based on other similar issues I've seen it seems like this has something to do with the interaction of reverse-mode autodiff, vmap, and the control flow used in odeint, but again, the issues I've seen raising these kinds of errors seem to have been solved?

Thanks!

[mattjj]

Thanks for raising this. It certainly looks like a bug! For that reason I made a corresponding bug issue in #8783. (Also I wanted to try out a button I just noticed on the GitHub UI...)

Let's follow up there!

[DanPuzzuoli]

Awesome, thanks!