message: Allow user to pass input_shape to split.
diff: @@ -90,7 +90,7 @@ class Sharding(object):
  tile_assignment_devices=list(flattened_devices)))
  @classmethod
  - def split(cls, tensor, split_dimension, num_devices):
  + def split(cls, tensor, split_dimension, num_devices, input_shape=None):
  """Returns a Sharding that splits a tensor across a dimension.
  This creates a Tiled attribute, similar to tile(), but easier to use for the
  @@ -100,11 +100,15 @@ class Sharding(object):
  tensor: A tf.Tensor to split.
  split_dimension: The dimension number to split.
  num_devices: The number of cores to split `tensor` over.
  + input_shape: The shape of the original tensor.
  Raises:
  ValueError: The tensor to split was smaller in the split dimension than
  the number of devices to split over.
  """
  + if input_shape:
  + shape = input_shape
  + else:
  shape = tensor.shape.as_list()
  if (shape[split_dimension] is not None and
  shape[split_dimension] < num_devices):
  @@ -221,7 +225,8 @@ def split(tensor,
  split_dimension,
  num_devices,
  assign_tuple_sharding=False,
  - use_sharding_op=False):
  + use_sharding_op=False,
  + input_shape=None):
  """Returns a tensor that is split along the given dimension.
  Args:
  @@ -230,10 +235,11 @@ def split(tensor,
  num_devices: The number of devices to partition the dimension.
  assign_tuple_sharding: If the sharding type should be a tuple.
  use_sharding_op: If true, adds a sharding op to set the sharding.
  + input_shape: The full shape of the input tensor.
  """
  if use_sharding_op:
  tensor = tf2xla.sharding(tensor)
  - Sharding.split(tensor, split_dimension, num_devices).apply_to_tensor(
  - tensor,
  - assign_tuple_sharding=assign_tuple_sharding)
  + Sharding.split(
  + tensor, split_dimension, num_devices, input_shape).apply_to_tensor(
  + tensor, assign_tuple_sharding=assign_tuple_sharding)
  return tensor
  