code: def test_grad_jit_metadata(self):
      @jax.jit
      def foo(x):
          return jnp.sin(x)
  
      hlo = module_to_string(
          jax.jit(jax.grad(foo)).lower(1.0).compiler_ir()
      )
      self.assertRegex(hlo, r'loc\(".*jvp\(jit\(foo\)\)/cos"')
      self.assertRegex(
          hlo, r'loc\(".*transpose\(jvp\(jit\(foo\)\)\)/mul"'
      )
  
file_path: jax/tests/metadata_test.py
imports: import io
  import unittest
  from absl.testing import absltest
  from jax._src import test_util
  import jax
  from jax._src import config
  from jax._src.lib.mlir import ir
  from jax import numpy
  from jax import config
label_desc: Use image node_id 3 for calling the MetadataTest obj's underlying member method code with example usage: obj.test_grad_jit_metadata() and returns: jnp