code: def cond(*args, **kwargs):
      """Equivalent to :func:`jax.lax.cond` but with Haiku state passed in/out.
  
      >>> true_fn = hk.nets.ResNet50(10)
      >>> false_fn = hk.Sequential([hk.Flatten(), hk.nets.MLP([300, 100, 10])])
      >>> x = jnp.ones([1, 224, 224, 3])
      >>> if hk.running_init():
      ...   # At `init` run both branches to create parameters everywhere.
      ...   true_fn(x)
      ...   out = false_fn(x)
      ... else:
      ...   # At `apply` conditionally call one of the modules.
      ...   i = jax.random.randint(hk.next_rng_key(), [], 0, 100)
      ...   out = hk.cond(i > 50, true_fn, false_fn, x)
  
      Args:
        pred: Boolean scalar type.
        true_fun: Function (A -> B), to be applied if ``pred`` is ``True``.
        false_fun: Function (A -> B), to be applied if ``pred`` is ``False``.
        operands: Operands (A) input to either branch depending on ``pred``. The
          type can be a scalar, array, or any pytree (nested Python tuple/list/dict)
          thereof.
  
      Returns:
        Value (B) of either ``true_fun(*operands)`` or ``false_fun(*operands)``,
        depending on the value of ``pred``. The type can be a scalar, array, or any
        pytree (nested Python tuple/list/dict) thereof.
      """
      # pylint: enable=g-doc-args
      if not base.inside_transform():
          raise ValueError(
              "hk.cond() should not be used outside of hk.transform(). "
              "Use jax.cond() instead."
          )
  
      try:
          bound_args = inspect.signature(_old_cond).bind(
              *args, **kwargs
          )
          (
              pred,
              true_operand,
              true_fun,
              false_operand,
              false_fun,
          ) = bound_args.args
          if not callable(true_fun) or not callable(false_fun):
              # Two operand new cond case: cond(pred, tf, ff, 1, 2)
              raise TypeError
      except TypeError:
          bound_args = inspect.signature(_new_cond).bind(
              *args, **kwargs
          )
          bound_args.apply_defaults()
          pred, true_fun, false_fun, *operands = bound_args.args
          operand = bound_args.kwargs["operand"]
          if operand is not SENTINEL:
              if operands:
                  raise ValueError(
                      "When the operand keyword argument is used you cannot "  # pylint: disable=raise-missing-from
                      "also pass operands positionally. Got "
                      f"operand={operand} and *operands={tuple(operands)}"
                  )
              operands = (operand,)
              del operand
      else:
          true_fun = lambda op, f=true_fun: f(op[0])
          false_fun = lambda op, f=false_fun: f(op[1])
          operands = ((true_operand, false_operand),)
  
      reserve_up_to_full_rng_block()
      stateful_branch_mem = _memoize_by_id(stateful_branch)
      state = internal_state()
      out, state = jax.lax.cond(
          pred,
          true_fun=stateful_branch_mem(true_fun),
          false_fun=stateful_branch_mem(false_fun),
          operand=(state, operands),
      )
      update_internal_state(state)
      return out
  
imports: import collections
  import collections.abc
  from collections.abc import Mapping, MutableMapping
  import functools
  import inspect
  from typing import Any, Callable, Optional, TypeVar
  from haiku._src import base
  import jax
  import jax.numpy
label_desc: Use image node_id 17 for calling a global function with example usage: cond() and returns: out