// RUN: tf-opt %s -tfl-prepare-quantize-dynamic-range="enable-custom-op-quantization=CustomTestOp=1" -tfl-quantize="enable-dynamic-range-quantization=true enable-custom-op-weight-only=CustomTestOp=false" -tfl-post-quantize="enable-no-side-effect=CustomTestOp=false" | FileCheck --check-prefix=NotPrune %s

// RUN: tf-opt %s -tfl-prepare-quantize-dynamic-range="enable-custom-op-quantization=CustomTestOp=1" -tfl-quantize="enable-dynamic-range-quantization=true enable-custom-op-weight-only=CustomTestOp=false" | FileCheck --check-prefix=CustomOpNotWeightOnly %s

// RUN: tf-opt %s -tfl-prepare-quantize-dynamic-range="enable-custom-op-quantization=CustomTestOp=1-3,CustomTestOp3=3" | FileCheck --check-prefix=CustomOp %s
// RUN: tf-opt %s -tfl-prepare-quantize-dynamic-range="min-elements-for-weights=4000 enable-custom-op-quantization=CustomTestOp=1-3,CustomTestOp3=3" | FileCheck --check-prefix=MinElement %s

  func.return %1 : tensor<i32>
}

func.func @UnsupportedOp(%arg0: tensor<i32>) -> tensor<i32> {
  // expected-error @+1 {{does not support subcomputations with tf/xla communication ops}}
  %0 = "tf.CustomTestOp"(%arg0) {config = "", config_proto = "", executor_type = "", f = @Callee} : (tensor<i32>) -> (tensor<i32>)
  func.return %0 : tensor<i32>
}

// -----

  LoadCustomOpTestModel();

  // The custom op is not hybrid, and the second input is a constant that can
  // be quantized.
  CustomOpMap custom_op_map;
  custom_op_map["CustomTestOp"] = {
      {1},   // quantizable_input_indices
      true,  // is_weight_only
  };

  flatbuffers::FlatBufferBuilder builder;
  auto status = QuantizeWeights(&builder, model_, 0, custom_op_map);

    func.return %0 : tensor<2xf32>
  }


  // CHECK-LABEL: unknown_op
  func.func @unknown_op(%arg0: tensor<2xf32>) -> tensor<2xf32> {
    // CHECK: tf.CustomTestOp
    %0 = "tf.CustomTestOp"(%arg0) : (tensor<2xf32>) -> tensor<2xf32>

    func.return %0 : tensor<2xf32>
  }

  // CHECK-LABEL: add_v2
  func.func @add_v2(%arg0: tensor<2xi32>) -> tensor<2xi32> {