// CHECK: %0 = "tfl.conv_2d"(%arg0, %arg1, %cst)
}

// CHECK-LABEL: fuse4DAddIntoConv2d
func.func @fuse4DAddIntoConv2d(%arg0: tensor<256x32x32x3xf32>, %arg1: tensor<2x3x3x3xf32>) -> tensor<256x32x32x2xf32> {
  %cst = arith.constant dense<[[[[1.0, 2.0]]]]> : tensor<1x1x1x2xf32>
  %cst_0 = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
  %0 = "tfl.conv_2d"(%arg0, %arg1, %cst_0) {
    dilation_h_factor = 1 : i32,

    save_options = None
    if has_func_alias:
      save_options = tensorflow.saved_model.SaveOptions(
          function_aliases={FUNC_ALIAS: model.conv2d}
      )
    saved_model_save.save(
        model,
        saved_model_path,
        signatures=model.conv2d.get_concrete_function(
            tensor_spec.TensorSpec(
                shape=input_shape, dtype=dtypes.float32, name='input_tensor'
            )

    dilation_width_factor = rewriter.getI32IntegerAttr(1);
  }

  StringAttr padding;
  if (!TFPaddingIsSameOrValid(op, &padding)) return failure();

  // TensorFlow Conv3D has no bias, optimization patterns will fuse Conv3D
  // with other ops can fill the bias.
  Value none = rewriter.create<TFL::NoValueOp>(
      op->getLoc(), rewriter.getNoneType(), rewriter.getUnitAttr());

  %2 = "tfl.pseudo_qconst"() {qtype = tensor<32x!quant.uniform<i32:f32, 1.7052092479439231E-4>>, value = dense<0> : tensor<32xi32>} : () -> tensor<32x!quant.uniform<i32:f32, 1.7052092479439231E-4>>

    func.return %add : tensor<*xf32>
  }

  func.func @composite_conv2d_with_relu6_fn(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32> attributes {tf_quant.composite_function} {

              (UpdateShapeWithAxis<-1> $qtype, $old_value))),
          [(CanUpdateShapeWithAxis<-1> $qtype, $old_value)]>;

// The axis is set to 0 because the transpose is from the legalization of
// tf.conv2d and the new channel axis is the first dimension.
def ReorderTransposeDequantQuantUsedByConv :
      Pat<(TF_TransposeOp:$old_value
              (TFL_DequantizeOp (TFL_QuantizeOp $input, $qtype)), $perm),

      if (!input_type || !input_type.hasStaticShape()) {
        return false;
      }
      total_count += input_type.getNumElements();
    }
    *count = total_count;
    return true;
  }

  // For conv2d/depthwise_conv/fully_connected ops.
  // This algorithm actually comes from TOCO tooling_util.cc
  static bool GetArithmeticCountForConvAndFullyconnectedOp(mlir::Operation* op,

      def body(self, x, w):
        z = nn_ops.conv2d(x, w, padding='SAME')
        return z, w

      @def_function.function(
          input_signature=[
              tensor_spec.TensorSpec(
                  shape=input_shape, dtype=dtypes.float32, name='input_tensor'
              )
          ]
      )
      def main(self, x):
        x1 = nn_ops.conv2d(x, self.w, padding='SAME')

    if ((quantization_method_ == tensorflow::quantization::QuantizationMethod::
                                     METHOD_DYNAMIC_RANGE_INT8) &&
        (function_name.contains("batch_matmul") ||
         function_name.contains("conv3d"))) {
      call_op->removeAttr(kQuantTraitAttrName);
    }

    // TODO(b/270906404): Support weight-only gather for uniform quantized opset
    // in PTQ mode
    if (target_opset_ == OpSet::UNIFORM_QUANTIZED &&

  %2 = "tfl.pseudo_qconst"() {qtype = tensor<32x!quant.uniform<i32:f32, 1.7052092479439231E-4>>, value = dense<0> : tensor<32xi32>} : () -> tensor<32x!quant.uniform<i32:f32, 1.7052092479439231E-4>>