Reshape x into a shape [num_devices, ...].

def _reshape_by_device_single(x, num_devices):
  """Reshape x into a shape [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:])