def METHOD_NAME(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)