%1 = "tf.PartitionedCall"(%arg0, %cst, %cst_0) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_1} : (tensor<1x2x2x3xf32>, tensor<2x2x3x2xf32>, tensor<2xf32>) -> tensor<*xf32> loc(callsite("test@conv"("Conv2D_1") at "QuantizationUnit(\12\08Conv2D_1\1a\04conv)"))

module {
func.func @main(%arg0: tensor<4x68x68x3xf32>, %arg1: tensor<5x5x3x8xf32>) -> tensor<4x64x64x8xf32> {
  %0 = "tf.Conv2D"(%arg0, %arg1) {padding = "VALID", strides = [1, 1, 1, 1]} : (tensor<4x68x68x3xf32>, tensor<5x5x3x8xf32>) -> tensor<4x64x64x8xf32>
  func.return %0 : tensor<4x64x64x8xf32>
}

   (IsInt8ElementType $filter),
   (IsConstTensor $filter),
   (IsInt32ElementType $conv),
   (HasStaticShapeConstraint $filter),
   (HasStaticShapeAtDimsConstraint<"3"> $input)],
  [], (addBenefit 10)>;

// Convert Conv2D with hybrid inputs (f32 activation/int8 weight) to XlaConv
def ConvertTFConv2DToXLAConvOpWeightOnly : Pat<
  (TF_Conv2DOp:$conv
    $input,
    (TF_MulOp (TF_CastOp (TF_IdentityOp $filter), $truncate1), $scale),

    %conv = "tf.Conv2D"(%arg0, %arg1) {attr_map = "0:strides,1:use_cudnn_on_gpu,2:padding,3:explicit_paddings,4:dilations", data_format = "NHWC", device = "", dilations = [1, 2, 2, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true} : (tensor<1x2x2x3xf32>, tensor<2x3x3x2xf32>) -> tensor<*xf32>
    return %conv : tensor<*xf32>
  }

  // CHECK-SAME: [1.300000e+01, 2.800000e+01], [1.500000e+01, 3.200000e+01], [1.700000e+01, 3.600000e+01]
  // CHECK: %[[CONV:.*]] = "tf.Conv2D"(%arg0, %[[CST]]) <{data_format = "NHWC", dilations = [1, 2, 3, 1], explicit_paddings = [], padding = "SAME", strides = [1, 4, 5, 1], use_cudnn_on_gpu = true}>
  // CHECK: return %[[CONV]] : tensor<1x28x23x2xf32>
}

// CHECK-LABEL: @notfuseMulIntoConv2d
// filter and multiply are not broadcastable

  // CHECK-NOT: "tf._FusedConv2D"
  %0 = "tf.Conv2D"(%arg2, %arg1) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}> : (tensor<8x32x32x3xf32>, tensor<1x1x3x128xf32>) -> tensor<*xf32>
  // The result of the conv must be the first input to BiasAdd to be fusable.

  %conv = "tf.Conv2D"(%dq_input, %dq_weight) {attr_map = "0:strides,1:use_cudnn_on_gpu,2:padding,3:explicit_paddings,4:dilations", data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<*xf32>

    key: "narrow_range"
    value {
      b: true
    }
  }
  attr {
    key: "num_bits"
    value {
      i: 4
    }
  }
}
node {
  name: "BoxPredictor_4/ClassPredictor/Conv2D"
  op: "Conv2D"
  input: "input"
  input: "BoxPredictor_4/ClassPredictor/weights_quant/FakeQuantWithMinMaxVarsPerChannel"
  attr {
    key: "T"
    value {
      type: DT_FLOAT
    }
  }
  attr {

	init.Append(as)
	return res
}

// Returns the data word (the second word) used to represent conv.X in
// an interface.
func dataWord(conv *ir.ConvExpr, init *ir.Nodes) ir.Node {
	pos, n := conv.Pos(), conv.X
	fromType := n.Type()

	// If it's a pointer, it is its own representation.
	if types.IsDirectIface(fromType) {
		return n
	}

        # One pure conv
        out = nn_ops.conv2d(
            out,
            self.conv_filters,
            strides=(1, 1, 2, 1),
            dilations=(1, 1, 1, 1),
            padding='SAME',
            data_format='NHWC',
        )

        # One fakequant attached conv
        if is_qat_model:
          out = array_ops.fake_quant_with_min_max_args(