reduce_op.getLoc(), reduce_op->getResult(0).getType(), operand,
          reduction_indices,
          /*keep_dim=*/rewriter.getBoolAttr(false));
      auto tf_argreduce_op = rewriter.create<ArgReduce>(
          reduce_op.getLoc(), reduce_op->getResult(1).getType(), operand,
          reduction_indices);

      rewriter.replaceOp(reduce_op, {tf_reduce_op, tf_argreduce_op});
    } else {

  if (!ranked_type) return UnrankedTensorType::get(builder->getI1Type());

  return RankedTensorType::get(ranked_type.getShape(), builder->getI1Type());
}

Type InferReductionOpType(Value input, Value reduction_indices,
                          BoolAttr keep_dims) {
  Type input_ty = input.getType();
  Type element_ty = getElementTypeOrSelf(input_ty);

  // Output type is unranked if input type is not ranked.

Type DeduceEqualCmpOpType(Builder *builder, Location loc, Value x, Value y,
                          BoolAttr incompatible_shape_error);

Type InferReductionOpType(Value input, Value reduction_indices,
                          BoolAttr keep_dims);

// Verifies that the given types are cast compatible. If not, emits appropriate
// error for the given op. If mask_one_dim is set to true, then the types are