code: def switch(index, branches, *operands):
      r"""Equivalent to :func:`jax.lax.switch` but with Haiku state passed in/out.
  
      Note that creating parameters inside a switch branch is not supported, as such
      at init time we recommend you unconditionally evaluate all branches of your
      switch and only use the switch at apply. For example:
  
      >>> experts = [hk.nets.MLP([300, 100, 10]) for _ in range(5)]
      >>> x = jnp.ones([1, 28 * 28])
      >>> if hk.running_init():
      ...   # During init unconditionally create params/state for all experts.
      ...   for expert in experts:
      ...     out = expert(x)
      ... else:
      ...   # During apply conditionally apply (and update) only one expert.
      ...   index = jax.random.randint(hk.next_rng_key(), [], 0, len(experts) - 1)
      ...   out = hk.switch(index, experts, x)
  
      Args:
        index: Integer scalar type, indicating which branch function to apply.
        branches: Sequence of functions (A -> B) to be applied based on index.
        operands: Operands (A) input to whichever branch is applied.
  
      Returns:
        Value (B) of branch(\*operands) for the branch that was selected based on
        index.
      """
      if not base.inside_transform():
          raise ValueError(
              "hk.switch() should not be used outside of hk.transform(). "
              "Use jax.switch() instead."
          )
  
      reserve_up_to_full_rng_block()
      stateful_branch_mem = _memoize_by_id(stateful_branch)
      state = internal_state()
      out, state = jax.lax.switch(
          index,
          tuple(python_map(stateful_branch_mem, branches)),
          (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 18 for calling a global function with example usage: switch(index, branches) and returns: out