from tensorflow.python.framework import ops


@ops.RegisterShape("SegmentMax")
@ops.RegisterShape("SegmentMean")
@ops.RegisterShape("SegmentMin")
@ops.RegisterShape("SegmentProd")
@ops.RegisterShape("SegmentSum")
def _SegmentReductionShape(op):
  """Common shape function for segment reduction ops."""
  data_shape = op.inputs[0].get_shape()
  segment_ids_shape = op.inputs[1].get_shape()
  segment_ids_shape.assert_has_rank(1)
  return [tensor_shape.TensorShape([None]).concatenate(data_shape[1:])]


@ops.RegisterShape("SparseSegmentMean")
@ops.RegisterShape("SparseSegmentSum")
def _SparseSegmentReductionShape(op):
  """Common shape function for sparse segment reduction ops."""
  data_shape = op.inputs[0].get_shape()
  indices_shape = op.inputs[1].get_shape()
  indices_shape.assert_has_rank(1)
  segment_ids_shape = op.inputs[2].get_shape()
  segment_ids_shape.assert_has_rank(1)
  indices_shape.assert_is_compatible_with(segment_ids_shape)
  return [tensor_shape.TensorShape([None]).concatenate(data_shape[1:])]


@ops.RegisterShape("SparseSegmentMeanGrad")
def _SparseSegmentMeanGradShape(op):
  """Shape function for the SparseSegmentMeanGrad op."""
