code: def test_grad_closure_with_vmap(self):
      # https://github.com/google/jax/issues/2718
      @jax.jit
      def experiment(x):
          def model(y, t):
              return -x * y
  
          history = odeint(model, 1.0, np.arange(0, 10, 0.1))
          return history[-1]
  
      gradfun = jax.value_and_grad(experiment)
      t = np.arange(0.0, 1.0, 0.01)
      h, g = jax.vmap(gradfun)(t)  # doesn't crash
      ans = h[11], g[11]
  
      expected_h = experiment(t[11])
      expected_g = (experiment(t[11] + 1e-5) - expected_h) / 1e-5
      expected = expected_h, expected_g
  
      self.assertAllClose(
          ans, expected, check_dtypes=False, atol=1e-2, rtol=1e-2
      )
  
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 11 for calling the ODETest obj's underlying member method code with example usage: obj.test_grad_closure_with_vmap() and returns: history, unknown