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 {

           [](mlir::OpBuilder& opb, mlir::Location loc, mlir::Value input,
              mlir::Value reduction_indices,
              bool keep_dims = false) -> mlir::Operation* {
             return opb
                 .create<mlir::TF::AnyOp>(loc, opb.getI1Type(), input,
                                          reduction_indices, keep_dims)
                 .getOperation();
           });

  // mlir::TF::ConstOp

  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.

           const char* raw_device_name = nullptr);

// Computes the sum of elements across dimensions of a tensor.
Status Sum(AbstractContext* ctx, AbstractTensorHandle* const input,
           AbstractTensorHandle* const reduction_indices,
           AbstractTensorHandle** output, bool keep_dims = false,
           const char* name = nullptr, const char* raw_device_name = nullptr);

// Returns x - y element-wise.

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