code: def expand_apply(f, axis=0):
      """Wraps f to temporarily add a size-1 axis to its inputs.
  
      Syntactic sugar for::
  
          ins = jax.tree_util.tree_map(lambda t: np.expand_dims(t, axis=axis), ins)
          out = f(ins)
          out = jax.tree_util.tree_map(lambda t: np.squeeze(t, axis=axis), out)
  
      This may be useful for applying a function built for ``[Time, Batch, ...]``
      arrays to a single timestep.
  
      Args:
        f: The callable to be applied to the expanded inputs.
        axis: Where to add the extra axis.
  
      Returns:
        f, wrapped as described above.
      """
      if axis not in [0, -1]:
          raise ValueError(
              "expand_apply currently only supports axis=0 or axis=-1."
          )
  
      @functools.wraps(f)
      def wrapper(*args, **kwargs):
          expand = lambda t: jnp.expand_dims(t, axis=axis)
          args = jax.tree_util.tree_map(expand, args)
          kwargs = jax.tree_util.tree_map(expand, kwargs)
          outputs = f(*args, **kwargs)
          return jax.tree_util.tree_map(
              lambda t: jnp.squeeze(t, axis=axis), outputs
          )
  
      return wrapper
  
imports: from collections.abc import Iterable, Sequence
  import functools
  from typing import Any, Callable, Optional
  from haiku._src import base
  from haiku._src import initializers
  from haiku._src import module
  from haiku._src import typing
  from haiku._src.typing import PRNGKey
  import jax
  from jax import lax
  import jax.numpy
  import numpy
label_desc: Use image node_id 7 for calling a global function with example usage: expand_apply(f, axis) and returns: wrapper, jax