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

def reshape_by_device(x, num_devices):
  """Reshape possibly nested x into a shape [num_devices, ...]."""
  return layers.nested_map(
      x, lambda x: _reshape_by_device_single(x, num_devices))