code: def _split_batch_dimension(
      new_batch: int, data: jnp.ndarray
  ) -> jnp.ndarray:
      """Splits the batch dimension and introduces new one with size `new_batch`.
  
      The result has two batch dimensions, first one of size `new_batch`, second one
      of size `data.shape[0]/new_batch`. It expects that `data.shape[0]` is
      divisible by `new_batch`.
  
      Args:
        new_batch: Dimension of outer batch dimension.
        data: jnp.ndarray to be reshaped.
  
      Returns:
        jnp.ndarray with extra batch dimension at start and updated second
        dimension.
      """
      # The first dimension will be used for allocating to a specific ensemble
      # member, and the second dimension is the parallelized batch dimension, and
      # the remaining dimensions are passed as-is to the wrapped network.
      # We use Fortan (F) order so that each input batch i is allocated to
      # ensemble member k = i % new_batch.
      return jnp.reshape(
          data, (new_batch, -1) + data.shape[1:], order="F"
      )
  
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 1 for calling a global function with example usage: _split_batch_dimension(new_batch, data) and returns: jnp