from tensorflow.python.framework import tensor_shape

def _SelectShape(op):
  # All three inputs must have the same shape.
  return [op.inputs[0].get_shape()
          .merge_with(op.inputs[1].get_shape())
          .merge_with(op.inputs[2].get_shape())]


@ops.RegisterShape("ArgMax")
@ops.RegisterShape("ArgMin")
def _ArgOpShape(op):
  """Common shape function for arg-reduction ops."""
  dimension_shape = op.inputs[1].get_shape()
  dimension_shape.assert_is_compatible_with(tensor_shape.scalar())
  input_shape = op.inputs[0].get_shape()
  if input_shape.ndims is None:
    return [tensor_shape.unknown_shape()]
  elif input_shape.ndims <= 1:
    return [tensor_shape.scalar()]

  dimension = tensor_util.ConstantValue(op.inputs[1])
  if dimension is None:
    return [tensor_shape.unknown_shape(ndims=input_shape.ndims - 1)]
  elif 0 <= dimension and dimension < input_shape.ndims:
    returned_shape = []
