code: def while_loop(
      cond_fn: Callable[[Scope, C], bool],
      body_fn: Callable[[Scope, C], C],
      scope: Scope,
      init: C,
      carry_variables: CollectionFilter = False,
      broadcast_variables: CollectionFilter = True,
      split_rngs: Mapping[PRNGSequenceFilter, bool] = {},
  ) -> C:
      """Lifted version of jax.lax.while_loop.
  
      The lifted scope is passed to `cond_fn` and `body_fn`.
      Broadcasted variables are immutable. The carry variable are
      mutable but cannot change shape and dtype.
      This also means you cannot initialize variables inside
      the body. Consider calling `body_fn` once manually before
      calling `while_loop` if variable initialization is required.
  
      Example::
  
        def f(scope, x):
          def cond_fn(scope, c):
            return scope.get_variable('state', 'acc') < 10
          def body_fn(scope, c):
            acc = scope.variable('state', 'acc')
            acc += 1
            y = scope.child(nn.dense)(c, c.shape[-1])
            return y
  
          c = x
          c = body_fn(scope, c)
          return lift.while_loop(cond_fn, body_fn, scope, (),
                                 carry_variables='state')
  
      Args:
        cond_fn: Should return True as long as the loop should continue.
        body_fn: The body of the while loop.
        scope: The scope(s) which should be lifted into the loop.
        init: The initial state passed to the loop
        carry_variables: collections that are carried through the loop
          and are therefore mutable (default: none).
        broadcast_variables: collections that are closed over and are
          therefore read-only (default: all collections)
        split_rngs: Split PRNG sequences will be different for each loop iterations.
          If split is False the PRNGs will be the same across iterations.
      Returns:
        The final state after executing the while loop.
      """
      rng_groups, rng_splits = _unzip2(split_rngs.items())
  
      def inner(scope_fn, repack_fn, variable_groups, rng_groups):
          carry_variables, broadcast_variables = variable_groups
  
          def make_loop_rngs(i):
              local_rng_groups = []
              for rng_group, rng_split in zip(rng_groups, rng_splits):
                  if rng_split:
                      rng_group = tree_map_rngs(
                          lambda rng: random.fold_in(rng, i), rng_group
                      )
                  local_rng_groups.append(rng_group)
              return local_rng_groups
  
          def cond_wrapper(c):
              i, carry_variables, carry = c
              scope = scope_fn(
                  (carry_variables, broadcast_variables),
                  make_loop_rngs(-i),
                  mutable_filter=False,
              )
              return cond_fn(scope, carry)
  
          def body_wrapper(c):
              i, carry_variables, carry = c
              scope = scope_fn(
                  (carry_variables, broadcast_variables),
                  make_loop_rngs(i),
              )
              carry = body_fn(scope, carry)
              (carry_variables,) = repack_fn(scope)
              return (i + 1, carry_variables, carry)
  
          c = (0, carry_variables, init)
          _, carry_variables, carry = jax.lax.while_loop(
              cond_wrapper, body_wrapper, c
          )
          return carry, (carry_variables,)
  
      return pack(
          inner,
          (carry_variables, broadcast_variables),
          (carry_variables,),
          rng_groups,
          name="while_loop",
      )(scope)
  
imports: import collections
  import dataclasses
  import functools
  from typing import Any, Callable, Dict, Generic, Iterable, List, Mapping, Optional, Sequence, Tuple, TypeVar, Union
  import warnings
  from flax import traceback_util
  import jax
  from jax import random
  from .None import axes_scan, meta
  from .frozen_dict import freeze, unfreeze
  from .scope import CollectionFilter, DenyList, Filter, PRNGSequenceFilter, Scope, group_collections, in_filter, intersect_filters, is_filter_empty, subtract_filters, union_filters
label_desc: Use image node_id 16 for calling a global function with example usage: while_loop(cond_fn, body_fn, scope, init, carry_variables, broadcast_variables, split_rngs) and returns: pack, carry,  carry_variables, local_rng_groups, cond_fn, unknown,  carry_variables,  carry