code: def _reshape_by_device_single(x, num_devices):
    
    x_shape = list(x.shape)
    batch_size = x_shape[0]
    batch_size_per_device = batch_size // num_devices
    # We require that num_devices divides batch_size evenly.
    if batch_size_per_device * num_devices != batch_size:
      logging.fatal(
          "We require that num_devices[%d] divides batch_size[%d] evenly.",
          num_devices, batch_size)
    # New shape.
    new_shape_prefix = [num_devices, batch_size_per_device]
    return np.reshape(x, new_shape_prefix + x_shape[1:])
docs: Reshape x into a shape [num_devices, ...].