code: def test_jvp_of_cond_transforms_name_stack(self):
      @jax.named_scope("foo")
      def f(x, y):
          @jax.named_scope("true")
          def true_fn(x):
              return x + 1
  
          @jax.named_scope("false")
          def false_fn(x):
              return x - 1
  
          return lax.cond(y, true_fn, false_fn, x)
  
      f_ = lambda x: jax.jit(f)(x, True)
      g = lambda x, t: jax.jvp(f_, (x,), (t,))
      jaxpr = jax.make_jaxpr(g)(jnp.arange(2.0), jnp.ones(2))
      jaxpr_param = "jaxpr"
      call_jaxpr = jaxpr.jaxpr.eqns[0].params[jaxpr_param]
      self.assertEqual(
          str(call_jaxpr.eqns[1].source_info.name_stack), "foo"
      )
      self.assertEqual(
          str(
              call_jaxpr.eqns[1]
              .params["branches"][0]
              .eqns[0]
              .source_info.name_stack
          ),
          "false",
      )
      self.assertEqual(
          str(
              call_jaxpr.eqns[1]
              .params["branches"][1]
              .eqns[0]
              .source_info.name_stack
          ),
          "true",
      )
  
      hlo_text = _get_hlo(g)(jnp.arange(2.0), jnp.ones(2))
      self.assertIn(
          "jvp(jit(f))/foo/cond/branch_0_fun/false/sub", hlo_text
      )
      self.assertIn(
          "jvp(jit(f))/foo/cond/branch_1_fun/true/add", 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 7 for calling the NameStackControlFlowTest obj's underlying member method code with example usage: obj.test_jvp_of_cond_transforms_name_stack() and returns: lax, unknown, unknown