from tensorflow.python.framework import tensor_util

  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."""
  input_shape = op.inputs[0].get_shape()
  indices_shape = op.inputs[1].get_shape().with_rank(1)
  unused_segment_ids_shape = op.inputs[2].get_shape().merge_with(indices_shape)
  unused_output_dim0_shape = op.inputs[3].get_shape().merge_with(
      tensor_shape.scalar())
  output_dim0 = tensor_util.ConstantValue(op.inputs[3])
  if output_dim0 is not None:
    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()
