docstr: A while with a single carry
docstr_tok: ["A","while","with","a","single","carry"]
code: def test_while_batched_cond(self):
      """A while with a single carry"""
  
      def product(x, y):
          # Equivalent to "x * y" implemented as:
          #      res = 0.
          #      for(i=0; i < y; i++)
          #         res += x
          return lax.while_loop(
              lambda idx_carry: idx_carry[0] < y,
              lambda idx_carry: (idx_carry[0] + 1, idx_carry[1] + x),
              (0, 0.0),
          )
  
      # We use vmap to compute result[i, j] = i * j
      xs = np.arange(4, dtype=np.int32)
      ys = np.arange(5, dtype=np.int32)
  
      def product_xs_y(xs, y):
          return jax.vmap(product, in_axes=(0, None))(xs, y)
  
      def product_xs_ys(xs, ys):
          return jax.vmap(product_xs_y, in_axes=(None, 0))(xs, ys)
  
      self.ConvertAndCompare(product_xs_ys, xs, ys)
  
imports: from absl.testing import absltest
  import jax
  import jax.lax
  import jax.numpy
  from jax._src import test_util
  import numpy
  from jax.experimental.jax2tf.tests import tf_test_util
  from jax import config
label_desc: Use image node_id 9 for calling the ControlFlowOpsTest obj's underlying member method code with example usage: obj.test_while_batched_cond() and returns: lax, jax, jax