code: def while_loop(
      cond_fn: Callable[[ModuleT, C], bool],
      body_fn: Callable[[ModuleT, C], C],
      mdl: ModuleT,
      init: C,
      carry_variables: lift.CollectionFilter = False,
      broadcast_variables: lift.CollectionFilter = True,
      split_rngs: Mapping[lift.PRNGSequenceFilter, bool] = FrozenDict(),
  ) -> 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::
  
        >>> import flax.linen as nn
        >>> import jax, jax.numpy as jnp
  
        >>> class WhileLoopExample(nn.Module):
        ...   @nn.compact
        ...   def __call__(self, x):
        ...     def cond_fn(mdl, c):
        ...       return mdl.variables['state']['acc'] < 10
        ...     def body_fn(mdl, c):
        ...       acc = mdl.variable('state', 'acc', lambda: jnp.array(0))
        ...       acc.value += 1
        ...       y = nn.Dense(c.shape[-1])(c)
        ...       return y
        ...     c = x
        ...     if self.is_mutable_collection('params'):
        ...       return body_fn(self, c)
        ...     else:
        ...       return nn.while_loop(cond_fn, body_fn, self, c,
        ...                             carry_variables='state')
  
        >>> k = jax.random.key(0)
        >>> x = jnp.ones((2, 2))
        >>> intial_vars = WhileLoopExample().init(k, x)
        >>> result, state = WhileLoopExample().apply(intial_vars, x, mutable=['state'])
  
      Args:
        cond_fn: Should return True as long as the loop should continue.
        body_fn: The body of the while loop.
        mdl: The Module 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.
      """
      return lift_direct_transform(
          lift.while_loop,
          (cond_fn, body_fn),
          mdl,
          init,
          carry_variables,
          broadcast_variables,
          split_rngs,
      )
  
imports: import dataclasses
  import functools
  import inspect
  from typing import Any, Callable, Dict, Iterable, Mapping, Optional, Sequence, Tuple, Type, TypeVar, Union
  from flax import errors, struct, traceback_util
  from flax.core import Scope, lift, meta
  from flax.core.frozen_dict import FrozenDict
  from flax.linen import module
  from flax.linen.module import Module, Variable, _derive_profiling_name, _get_unbound_fn, wrap_method_once
  import jax
label_desc: Use image node_id 21 for calling a global function with example usage: while_loop(cond_fn, body_fn, mdl, init, carry_variables, broadcast_variables, split_rngs) and returns: lift_direct_transform