// CHECK: func @main(%arg0: tensor<?x19x19x3xf32>) -> tensor<?x9x9x4xf32>
func.func @main(%arg0: tensor<?x19x19x3xf32>) -> tensor<?x9x9x4xf32> {
  %cst = arith.constant dense<1.0> : tensor<4xf32>
  %cst_3 = arith.constant dense<2.0> : tensor<4x3x3x3xf32>

class PrepareDynamicRangeQuantizableOp
    : public OpRewritePattern<arith::ConstantOp> {
 public:
  explicit PrepareDynamicRangeQuantizableOp(
      MLIRContext* context, const quant::QuantizationSpecs& quant_specs,
      llvm::SetVector<Operation*>* const visited_nonquantizable_ops)
      : OpRewritePattern<arith::ConstantOp>(context),
        visited_nonquantizable_ops_(visited_nonquantizable_ops),

// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1xf32>) -> tensor<1xf32>
attributes  {tf.entry_function = {inputs = "input", outputs = "output_node"}} {
  %0 = arith.constant dense<2.000000e+00> : tensor<f32>
  %1 = arith.constant dense<1.000000e+00> : tensor<f32>
  %2 = "tf.AddV2"(%arg0, %1) {T = "tfdtype$DT_FLOAT", device = "", name = "StatefulPartitionedCall/add"} : (tensor<1xf32>, tensor<f32>) -> tensor<1xf32>

limitations under the License.
==============================================================================*/

#include <memory>
#include <utility>

#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/IR/PatternMatch.h"  // from @llvm-project
#include "mlir/Pass/PassRegistry.h"  // from @llvm-project

limitations under the License.
==============================================================================*/

#include "tensorflow/compiler/mlir/tf2xla/transforms/xla_legalize_targets.h"

#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
#include "mlir/Dialect/Shape/IR/Shape.h"  // from @llvm-project
#include "mlir/Dialect/Tensor/IR/Tensor.h"  // from @llvm-project

func.func @mixed_no_constant_folding() -> (tensor<f32>) {
  // CHECK:      %[[CST0:.+]] = arith.constant dense<0.000000e+00>
  // CHECK-NEXT: %[[CST1:.+]] = arith.constant dense<0.000000e+00>
  // CHECK-NEXT: "vhlo.add_v1"(%[[CST0]], %[[CST1]]) : (tensor<f32>, tensor<f32>) -> tensor<f32>
  %cst_0 = arith.constant dense<0.000000e+00> : tensor<f32>
  %cst_1 = arith.constant dense<0.000000e+00> : tensor<f32>
  %0 = stablehlo.add %cst_0, %cst_1 : tensor<f32>

module {
func.func @tfInplaceUpdate(%arg0: tensor<2x1x2xf32>) -> tensor<2x1x2xf32> {
  %1 = arith.constant dense<1> : tensor<1xi32>
  %2 = arith.constant dense<2.0> : tensor<1x1x2xf32>
  %3 = "tf.InplaceUpdate"(%arg0, %1, %2) {device = ""}
    : (tensor<2x1x2xf32>, tensor<1xi32>, tensor<1x1x2xf32>) -> tensor<2x1x2xf32>

// CHECK-DAG:       [[VAL_8:%.*]] = arith.constant dense<0.000000e+00> : tensor<16x16xf32>
// CHECK-DAG:       [[VAL_9:%.*]] = arith.constant dense<0.000000e+00> : tensor<16x8xf32>
// CHECK-DAG:       [[VAL_10:%.*]] = arith.constant dense<0.000000e+00> : tensor<16xf32>
// CHECK-DAG:       [[VAL_11:%.*]] = arith.constant dense<0.000000e+00> : tensor<1x16xf32>
// CHECK-DAG:       [[VAL_12:%.*]] = arith.constant dense<0.000000e+00> : tensor<8x16xf32>

// RUN: odml-to-stablehlo-opt %s -tf-stablehlo | FileCheck %s

module attributes {tf.versions = {producer = 888 : i32}} {

func.func @tfInplaceUpdate(%arg0: tensor<2x1x2xf32>) -> tensor<2x1x2xf32> {
  %1 = arith.constant dense<1> : tensor<1xi32>
  %2 = arith.constant dense<2.0> : tensor<1x1x2xf32>
  %3 = "tf.InplaceUpdate"(%arg0, %1, %2) {device = ""}
    : (tensor<2x1x2xf32>, tensor<1xi32>, tensor<1x1x2xf32>) -> tensor<2x1x2xf32>

    }
  });
}

arith::ConstantOp QuantizationDriver::DuplicateConstantOpIfNeeded(
    arith::ConstantOp op, Operation* target_op, const int operand_index) {
  if (op.getResult().hasOneUse()) {
    return op;
  }
  OpBuilder builder(op->getContext());
  builder.setInsertionPointAfter(op);
  arith::ConstantOp new_op = cast<arith::ConstantOp>(builder.clone(*op));