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 {

  // Running example:
  // input: [[5, 5, 5],
  //         [1, 2, -3]]
  // reduction_indices: [1]
  auto input = op.input(0);
  auto reduction_indices = op.input(1);

  // [2, 3]
  auto input_shape = Shape(scope, input);

  // [2, 1]
  auto output_shape_kept_dims =
      ReducedShapeHelper(scope, input_shape, reduction_indices);

  // for op=min (say)
  // output = [5, -3]
  // y = [[5],

           [](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.

      CompileTimeConstantOperand<TF::MaxOp, 1>,  // $reduction_indices
      // $ksize, $strides
      CompileTimeConstantOperand<TF::MaxPoolGradGradV2Op, 3, 4>,
      // $ksize, $strides
      CompileTimeConstantOperand<TF::MaxPoolGradV2Op, 2, 3>,
      CompileTimeConstantOperand<TF::MaxPoolV2Op, 1, 2>,   // $ksize, $strides
      CompileTimeConstantOperand<TF::MeanOp, 1>,           // $reduction_indices

//   in `axis`. If `keep_dims` is true, the reduced dimensions are retained with
//   length 1.
Status Sum(AbstractContext* ctx, AbstractTensorHandle* const input,
           AbstractTensorHandle* const reduction_indices,
           AbstractTensorHandle** output, bool keep_dims, const char* name,
           const char* raw_device_name) {
  AbstractOperationPtr op_ptr(ctx->CreateOperation());

def LegalizeAny : Pat<
  (TF_AnyOp $input, $reduction_indices, $keep_dims),
  (TFL_ReduceAnyOp $input, (CreateTFCastToInt32Op $reduction_indices),
                   $keep_dims)>;

def LegalizeAll : Pat<
  (TF_AllOp $input, $reduction_indices, $keep_dims),
  (TFL_ReduceAllOp $input, (CreateTFCastToInt32Op $reduction_indices),
                   $keep_dims)>;

           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

//===----------------------------------------------------------------------===//

// Size(x) = Prod(Shape(x), reduction_indices=0, keep_dims=false)
def LowerSizeOp : Pat<
  (TF_SizeOp:$res $arg),
  (TF_ProdOp
    (CreateTFShapeOp
      $res,
      $arg,
      (IsI32 $res)
    ),
    /*reduction_indices=*/
    (TF_ConstOp
      (GetScalarOfType<0> $res)
    ),
    /*keep_dims=*/
    ConstBoolAttrFalse