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