code: def apply_all(
      base_apply: Callable[[networks.Params, Any], Any],
      params: networks.Params,
      *args,
      **kwargs,
  ) -> Any:
      """Pass the input to all ensemble members.
  
      Inputs can either have a batch dimension which will get implicitly vmapped
      over, or can be a single vector which will get sent to all ensemble members.
      e.g. [<inputs_dims>] or [batch_size, <input_dims>].
  
      Args:
        base_apply: Base network `apply` function that will be used for averaging.
          NOTE -- This will not work with mutable/stateful apply functions. --
        params: Model parameters.  Number of networks is deduced from this.
        *args: Allows for arbitrary call signatures for `base_apply`.
        **kwargs: Allows for arbitrary call signatures for `base_apply`.
  
      Returns:
        pytree of the resulting output of passing input to all ensemble members.
        Output shape will be [num_members, batch_size, <network output dims>].
      """
      # `num_networks` is the size of the batch dimension in `params`.
      num_networks = jax.tree_util.tree_leaves(params)[0].shape[0]
  
      args = jax.tree_map(
          functools.partial(_repeat_n, num_networks), args
      )
      kwargs = jax.tree_map(
          functools.partial(_repeat_n, num_networks), kwargs
      )
      # `out` is of shape `(num_networks, batch_size, <remaining dimensions>)`.
      return jax.vmap(base_apply)(params, *args, **kwargs)
  
file_path: acme/acme/agents/jax/mbop/ensemble.py
imports: import functools
  from typing import Any, Callable
  from acme.jax import networks
  import jax
  import jax.numpy
label_desc: Use image node_id 5 for calling a global function with example usage: apply_all(base_apply, params) and returns: jax