// CHECK: %[[scast:.*]] = "quantfork.scast"(%[[conv]]
// CHECK: %[[fcast:.*]] = "tf.Cast"(%[[scast]]) <{Truncate = false}> : (tensor<*xi8>) -> tensor<*xf32>
// CHECK: %[[avgpool_f32:.*]] = "tf.AvgPool"(%[[fcast]])
// CHECK-SAME: (tensor<*xf32>) -> tensor<*xf32>
// CHECK: %[[round:.*]] = "tf.Round"(%[[avgpool_f32]])
// CHECK: %[[icast:.*]] = "tf.Cast"(%[[round]]) <{Truncate = false}> : (tensor<*xf32>) -> tensor<*xi8>

    %5 = "quantfork.dcast"(%4) : (tensor<2x3x1024x3x!quant.uniform<i8<-127:127>:f32, 4.000000e-01>>) -> tensor<2x3x1024x3xf32>
    %6 = "quantfork.qcast"(%arg0) {volatile} : (tensor<2x3x1x1024xf32>) -> tensor<2x3x1x1024x!quant.uniform<i8:f32, 5.000000e-01:2>>

    // CHECK: %[[DQ:.*]] = "quantfork.dcast"(%[[RESHAPE]]) : (tensor<3x1x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<3x1xf32>
    // CHECK: return %[[DQ]]
    %0 = "quantfork.qcast"(%arg0) {volatile} : (tensor<1x2xf32>) -> tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03>>
    %1 = "quantfork.dcast"(%0) : (tensor<1x2x!quant.uniform<i8:f32, 5.000000e-03>>) -> tensor<1x2xf32>

// CHECK: %[[CST:.+]] = "tf.Const"() <{value = dense<3.000000e-01> : tensor<2x3x3x2xf32>}> : () -> tensor<2x3x3x2xf32>
// CHECK: %[[Q:.+]] = "quantfork.qcast"(%[[CST]]) : (tensor<2x3x3x2xf32>) -> tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, 0.0023622048182750312>>
// CHECK: %[[DQ:.+]] = "quantfork.dcast"(%[[Q]]) : (tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, 0.0023622048182750312>>) -> tensor<2x3x3x2xf32>

  std::vector<Value> sliced_lhs =
      sliceInput(input_lhs, bcast.x_batch_size(), loc, rewriter);
  std::vector<Value> sliced_rhs =
      sliceInput(input_rhs, bcast.y_batch_size(), loc, rewriter);

  // Compute (single batch) MatMul for each output batch.
  std::vector<Value> matmuls;
  matmuls.reserve(bcast.output_batch_size());
  for (int batch_idx : llvm::seq<int>(0, bcast.output_batch_size())) {
    int lhs_batch_idx, rhs_batch_idx;

  ];
}

def PrepareQuantizePass : Pass<"stablehlo-prepare-quantize", "mlir::ModuleOp"> {
  let summary = "Prepare StableHLO dialect for static range quantization by converting quantfork.stats into quantfork.qcast and dcast ops.";
  let options = [
    Option<"enable_per_channel_quantized_weight_",
        "enable-per-channel-quantized-weight",
        "bool", /*default=*/"true",
        "Whether to enable per-channel quantized weights.">,

      %1 = "quantfork.qcast"(%0) {volatile} : (tensor<1x3xf32>) -> tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
      %2 = "quantfork.dcast"(%1) : (tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<1x3xf32>
      %3 = stablehlo.reshape %2 : (tensor<1x3xf32>) -> tensor<3x1xf32>

  %0 = "quantfork.qcast"(%cst_1) : (tensor<2x3x3x2xf32>) -> tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {0.003937007874015748,0.003937007874015748}>>
  %1 = "quantfork.dcast"(%0) : (tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {0.003937007874015748,0.003937007874015748}>>) -> tensor<2x3x3x2xf32>

// TODO(b/204265523): Removes this pass after the exporting MLIR to SavedModel
// path is available.
std::unique_ptr<OperationPass<ModuleOp>> CreateInsertMainFunctionPass();

// Converts FakeQuant ops to quant.qcast and quant.dcast (QDQ) pairs.
std::unique_ptr<OperationPass<func::FuncOp>> CreateConvertFakeQuantToQdqPass();

// Lifts the quantizable spots as composite functions.
std::unique_ptr<OperationPass<ModuleOp>>

      ConstantOp new_const = rewriter.create<ConstantOp>(
          op->getLoc(), new_result_type, new_value_attr);
      ConvertOp dcast =
          rewriter.create<ConvertOp>(op->getLoc(), old_result_type, new_const);
      // replace all convert ops with dq op.
      convert_op->replaceAllUsesWith(dcast);
      // Return without scanning for the next ConvertOp as only one ConvertOp is
      // connected to all quantizable ops.