(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')

  %dq_bias = "quantfork.dcast"(%q_bias) : (tensor<2x!quant.uniform<i32:f32, 0.044022349891595126>>) -> tensor<2xf32>

    func.return %arg0 : tensor<*xi32>
  }

  // Test conv2d inferReturnTypes can infer some information when input or
  // filter does not have fully static shape.

  // CHECK-LABEL: func @conv2d_unranked_input_and_filter
  func.func @conv2d_unranked_input_and_filter(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32> {
    // CHECK: "tf.Conv2D"
    // CHECK-SAME: -> tensor<?x?x?x?xf32>

  %0 = "tf.Conv2D"(%output, %cst) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 2, 2, 1], use_cudnn_on_gpu = true}> {attr_map = "0:strides,1:use_cudnn_on_gpu,2:padding,3:explicit_paddings,4:dilations", device...

def MultiplyFakeQuantValue : NativeCodeCall<
  "MultiplyFakeQuantValue($_builder, $_loc, $0...)">;

// Convert AddV2Op following an AffineOp to BiasAddOp.
// For Conv3D, even though the Conv3D op has "NDHWC" data format, the BiasAdd
// will still has the data format of "NHWC".
def ConvertAddToBiasAdd : Pat<
  (TF_AddV2Op
    (SupportedAffineOpMatcher $conv_out, $input, $weight),

  %dq_bias = "quantfork.dcast"(%q_bias) : (tensor<2x!quant.uniform<i32:f32, 0.044022349891595126>>) -> tensor<2xf32>

  %input_scale = "tf.Const"() {value = dense<1.0> : tensor<f32>} : () -> tensor<f32>
  %side_input_scale = "tf.Const"() {value = dense<2.0> : tensor<f32>} : () -> tensor<f32>
  %conv2d = "tf._FusedConv2D"(%input, %filter, %bias, %act, %input_scale, %side_input_scale) {

      %outputs_6, %control_7 = tf_executor.island wraps "tf.Const"() {device = "", value = dense<[-1, 16384]> : tensor<2xi32>} : () -> tensor<2xi32>
      %outputs_8, %control_9 = tf_executor.island wraps "tf.Conv2D"(%arg0, %outputs_0) {data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 2, 2, 1], use_cudnn_on_gpu = true} : (tensor<16x224x224x3xf32>, tensor<*xf32>) -> tensor<16x112x112x?xf32>