code: def _while_loop_jax(
      cond,
      body,
      loop_vars,  # pylint: disable=redefined-outer-name
      shape_invariants=None,
      parallel_iterations=10,  # pylint: disable=unused-argument
      back_prop=True,
      swap_memory=False,  # pylint: disable=unused-argument
      maximum_iterations=None,
      name=None,
  ):  # pylint: disable=unused-argument
      """Jax implementation of `tf.while_loop`."""
      from jax import lax  # pylint: disable=g-import-not-at-top
  
      pack_body = lambda x: nest.pack_sequence_as(
          loop_vars, nest.flatten(x)
      )
  
      if maximum_iterations is None:
  
          def override_body_fn(args):
              return pack_body(body(*args))
  
          def override_cond_fn(args):
              return cond(*args)
  
          return lax.while_loop(
              override_cond_fn, override_body_fn, loop_vars
          )
      elif back_prop:
  
          def override_body_fn(args, _):
              c = cond(*args)
              sc = ops.get_static_value(c)
              if sc is None:
                  args = lax.cond(
                      c,
                      args,
                      lambda args: pack_body(body(*args)),
                      args,
                      lambda args: args,
                  )
              elif sc:
                  args = pack_body(body(*args))
              return args, ()
  
          loop_vars, _ = lax.scan(
              override_body_fn,
              loop_vars,
              xs=None,
              length=maximum_iterations,
          )
          return loop_vars
      else:
  
          def override_body_fn(args):
              i, args = args
              return i + 1, pack_body(body(*args))
  
          def override_cond_fn(args):
              i, args = args
              return cond(*args) & (i < maximum_iterations)
  
          return lax.while_loop(
              override_cond_fn,
              override_body_fn,
              (np.array(0), loop_vars),
          )[1]
  
file_path: probability/tensorflow_probability/python/internal/backend/numpy/control_flow.py
imports: import collections
  import functools
  import numpy
  from tensorflow_probability.python.internal.backend.numpy import _utils
  from tensorflow_probability.python.internal.backend.numpy import dtype
  from tensorflow_probability.python.internal.backend.numpy import nest
  from tensorflow_probability.python.internal.backend.numpy import ops
label_desc: Use image node_id 5 for calling a global function with example usage: _while_loop_jax(cond, body, loop_vars, shape_invariants, parallel_iterations, back_prop, swap_memory, maximum_iterations, name) and returns: lax, pack_body, cond, loop_vars, lax, args,  , unknown,  pack_body, unknown