from tensorflow.python.framework import ops

    dim0 = output_dim0[0]
  else:
    dim0 = None
  return [tensor_shape.TensorShape([dim0]).concatenate(input_shape[1:])]


@ops.RegisterShape("UnsortedSegmentSum")
def _UnsortedSegmentSumShape(op):
  """Shape function for UnsortedSegmentSum."""
  data_shape = op.inputs[0].get_shape()
  segment_ids_shape = op.inputs[1].get_shape()
  mid = segment_ids_shape.ndims
