from tensorflow.python.framework import tensor_shape

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
  if mid is None:
    return [tensor_shape.unknown_shape()]
  else:
    num_segments = tensor_util.ConstantValue(op.inputs[2])
    return [tensor_shape.TensorShape([num_segments]).concatenate(
        data_shape[mid:])]


@ops.RegisterShape("LinSpace")
def _LinspaceShape(op):
  num = tensor_util.ConstantValue(op.inputs[2])
  return [tensor_shape.vector(num)]
