code: def test_vmap_of_cond_should_transform_name_stack(self):
      @jax.named_scope("foo")
      @functools.partial(jax.vmap, in_axes=(0, None))
      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)
  
      jaxpr = jax.make_jaxpr(f)(jnp.arange(2), True)
      for eqn in jaxpr.eqns:
          self.assertIn("foo", str(eqn.source_info.name_stack))
          if eqn.primitive is lax.cond_p:
              self.assertEqual(
                  str(
                      eqn.params["branches"][0]
                      .eqns[0]
                      .source_info.name_stack
                  ),
                  "false",
              )
              self.assertEqual(
                  str(
                      eqn.params["branches"][1]
                      .eqns[0]
                      .source_info.name_stack
                  ),
                  "true",
              )
  
      hlo_text = _get_hlo(f)(jnp.arange(2.0), True)
      self.assertIn("foo/vmap(cond)/branch_0_fun/false/sub", hlo_text)
      self.assertIn("foo/vmap(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 6 for calling the NameStackControlFlowTest obj's underlying member method code with example usage: obj.test_vmap_of_cond_should_transform_name_stack() and returns: lax, unknown, unknown