code: def test_odeint_vmap_grad(self):
      # https://github.com/google/jax/issues/2531
  
      def dx_dt(x, *args):
          return 0.1 * x
  
      def f(x, y):
          y0 = jnp.array([x, y])
          t = jnp.array([0.0, 5.0])
          y = odeint(dx_dt, y0, t)
          return y[-1].sum()
  
      def g(x):
          # Two initial values for the ODE
          y0_arr = jnp.array([[x, 0.1], [x, 0.2]])
  
          # Run ODE twice
          t = jnp.array([0.0, 5.0])
          y = jax.vmap(lambda y0: odeint(dx_dt, y0, t))(y0_arr)
          return y[:, -1].sum()
  
      ans = jax.grad(g)(2.0)  # don't crash
      expected = jax.grad(f, 0)(2.0, 0.1) + jax.grad(f, 0)(2.0, 0.2)
  
      atol = {jnp.float32: 1e-5, jnp.float64: 5e-15}
      rtol = {jnp.float32: 1e-5, jnp.float64: 2e-15}
      self.assertAllClose(
          ans, expected, check_dtypes=False, atol=atol, rtol=rtol
      )
  
imports: from functools import partial
  from absl.testing import absltest
  import numpy
  import jax
  from jax._src import test_util
  import jax.numpy
  from jax.experimental.ode import odeint
  from jax.tree_util import tree_map
  import scipy.integrate
  from jax import config
label_desc: Use image node_id 8 for calling the ODETest obj's underlying member method code with example usage: obj.test_odeint_vmap_grad() and returns: unknown, y, y