code: def test_jvp_should_apply_to_call_jaxpr(self):
      @jax.jit
      def f(x):
          with jax.named_scope("bar"):
              with jax.named_scope("baz"):
                  return jnp.square(x)
  
      g = jax.named_scope("foo")(lambda x, t: jax.jvp(f, (x,), (t,)))
      jaxpr = jax.make_jaxpr(g)(1.0, 1.0).jaxpr
      jaxpr_param = "jaxpr"
      self.assertEqual(str(jaxpr.eqns[0].source_info.name_stack), "foo")
      self.assertEqual(
          str(
              jaxpr.eqns[0]
              .params[jaxpr_param]
              .eqns[0]
              .source_info.name_stack
          ),
          "bar/baz",
      )
  
      hlo_text = _get_hlo(g)(1.0, 1.0)
      self.assertIn("foo/jvp(jit(f))/bar/baz/mul", hlo_text)
  
imports: import functools
  from absl.testing import absltest
  import jax
  import jax.numpy
  from jax._src import core
  from jax import lax
  from jax._src.pjit import pjit
  from jax._src import linear_util
  from jax import config
  from jax._src import test_util
  from jax._src.lib import xla_client
label_desc: Use image node_id 5 for calling the NameStackTransformationTest obj's underlying member method code with example usage: obj.test_jvp_should_apply_to_call_jaxpr() and returns: jnp