code: def apply_mean(
      base_apply: Callable[[networks.Params, Any], Any],
      params: networks.Params,
      *args,
      **kwargs,
  ) -> Any:
      """Calculates the mean over all ensemble members for each batch element.
  
      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 average over all ensembles for each element.
        Output shape will be [batch_size, <network output_dims>]
      """
      out = apply_all(base_apply, params, *args, **kwargs)
      return jax.tree_map(functools.partial(jnp.mean, axis=0), out)
  
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 6 for calling a global function with example usage: apply_mean(base_apply, params) and returns: jax