code: def apply_round_robin(
      base_apply: Callable[[networks.Params, Any], Any],
      params: networks.Params,
      *args,
      **kwargs,
  ) -> Any:
      """Passes the input in a round-robin manner.
  
      The round-robin application means that each element of the input batch will
      be passed through a single ensemble member in a deterministic round-robin
      manner, i.e. element_i -> member_k where k = i % num_networks.
  
      It expects that:
      * `base_apply(member_params, *member_args, **member_kwargs)` is a valid call,
         where:
        * `member_params.shape = params.shape[1:]`
        * `member_args` and `member_kwargs` have the same structure as `args` and
          `kwargs`.
      * `params[k]` contains the params of the k-th member of the ensemble.
      * All jax arrays in `args` and `kwargs` have a batch dimension at axis 0 of
        the same size, which is divisible by `params.shape[0]`.
  
      Args:
        base_apply: Base network `apply` function that will be used for round-robin.
          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 round-robin application.
        Output shape will be [initial_batch_size, <remaining dimensions>].
      """
      # `num_networks` is the size of the batch dimension in `params`.
      num_networks = jax.tree_util.tree_leaves(params)[0].shape[0]
  
      # Reshape args and kwargs for the round-robin:
      args = jax.tree_map(
          functools.partial(_split_batch_dimension, num_networks), args
      )
      kwargs = jax.tree_map(
          functools.partial(_split_batch_dimension, num_networks),
          kwargs,
      )
      # `out.shape` is `(num_networks, initial_batch_size/num_networks, ...)
      out = jax.vmap(base_apply)(params, *args, **kwargs)
      # Reshape to [initial_batch_size, <remaining dimensions>]. Using the 'F' order
      # forces the original values to the last dimension.
      return jax.tree_map(
          lambda x: x.reshape((-1,) + x.shape[2:], order="F"), 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 4 for calling a global function with example usage: apply_round_robin(base_apply, params) and returns: jax