// RUN: tf-quant-opt %s -split-input-file -quant-preprocess-op='target-opset=XLA quantization-method=weight_only enable-per-channel-quantization=false' | FileCheck --check-prefix PerTensor %s
// RUN: tf-quant-opt %s -split-input-file -quant-preprocess-op='target-opset=XLA quantization-method=weight_only enable-per-channel-quantization=true' | FileCheck --check-prefix PerChannel %s

module {
  // For XLA weight-only per-channel depthwise convolution, tensor shape should have

// RUN: tf-quant-opt %s -quant-insert-quantized-functions='quantization-method=weight_only target-opset=XLA' | FileCheck %s

// Empty module
module {
  func.func @simple_fn(%arg0: tensor<*xf32>) -> tensor<*xf32> {
    func.return %arg0 : tensor<*xf32>
  }
}

// CHECK-NOT: func private @internal_dequantize_f32
// CHECK-NOT: func private @internal_conv3d_fn
// CHECK-NOT: func private @internal_batch_matmul_fn

          clEnumValN(tensorflow::quantization::QuantizationMethod::
                         METHOD_STATIC_RANGE_WEIGHT_ONLY_INT8,
                     "weight_only", "Post-training weight_only quantizaiton"))};
};

class CheckQuantizableOps
    : public mlir::OpRewritePattern<TF::PartitionedCallOp> {
 public:
  explicit CheckQuantizableOps(MLIRContext* context,

          clEnumValN(tensorflow::quantization::QuantizationMethod::
                         METHOD_STATIC_RANGE_WEIGHT_ONLY_INT8,
                     "weight_only", "Post-training weight_only quantizaiton"))};

  Option<OpSet> op_set_{
      *this, "target-opset", llvm::cl::init(OpSet::TF),
      llvm::cl::desc("Choose target opset."),
      llvm::cl::values(

      break;
    // Note: This is weight-only quantization by default, but with the legacy
    // flag "--force_dynamic_range_in_kernel", a DRQ behavior will be forced
    // in the kernel.
    case PresetQuantizationMethod::WEIGHT_ONLY:
      weight_component = custom_method.add_quantization_component_spec();
      SetQuantizationComponentSpec(weight_component,
                                   QuantizationComponentSpec::COMPONENT_WEIGHT,

    // Apply default weight-only quantization. Weights are quantized during
    // conversion, then dequantized during inference.
    // Activation: f32, Weight: qi8, Bias: f32
    WEIGHT_ONLY = 1;

    // Apply default dynamic range quantization. Quantized tensor value's
    // ranges are determined during graph runtime.
    // Activation: f32, Weight: qi8, Bias: f32

  // is:
  // 1) STATIC_RANGE
  //    - COMPONENT_ACTIVATION: INT_8
  //    - COMPONENT_WEIGHT: INT_8
  //    - COMPONENT_BIAS: INT_32
  // 2) WEIGHT_ONLY
  //    - COMPONENT_WEIGHT: INT_8
  // 3) DYNAMIC_RANGE
  //    - COMPONENT_ACTIVATION: INT_8
  //    - COMPONENT_WEIGHT: INT_8
  //    - COMPONENT_BIAS: INT_32

          custom_op_map[node_name].quantizable_input_indices.push_back(
              std::stoi(cur_index));
        }
        break;
      }
      case CustomOpUpdateOptions::kWeightOnly:
        custom_op_map[node_name].is_weight_only =
            GetBooleanSpecs(node_specification);
        break;
      case CustomOpUpdateOptions::kNoSideEffect:
        custom_op_map[node_name].no_side_effect =