func.func @convolution_add_add(
    %lhs: tensor<?x3x2x1xi8>, %rhs: tensor<2x1x1x1xi8>,
    %zp_offset: tensor<?x2x2x1xi32>, %bias: tensor<1xi32>
  ) -> tensor<?x2x2x1xi32> {
  // CHECK-DAG: %[[conv:.*]] = mhlo.convolution
  // CHECK-DAG: %[[combined:.*]] = chlo.broadcast_add %[[zp_offset:.*]], %[[bias:.*]]
  // CHECK-DAG: %[[result:.*]] = chlo.broadcast_add %[[conv]], %[[combined]]
  // CHECK: return %[[result]]

          mlir::cast<quant::QuantizedType>(non_bias_type.getElementType());
      non_bias_types.push_back(non_bias_ele_type);
    } else {
      // The non-bias hasn't been quantized, let's skip this bias.
      break;
    }
  }
  // The non-bias hasn't been quantized, let's skip this bias.
  if (non_bias_types.size() != non_biases.size()) return {};

  return func(/*op_types=*/non_bias_types, /*adjusted_quant_dim=*/-1,

  [(IsNotInLiftedFunc $res), (IsStableHLOConstantOp $bias)], [], (addBenefit 5)>;

def LiftDotGeneralWithBias : Pat<
  (StableHLO_AddOp:$res
    (StableHLO_DotGeneralOp
        $lhs, $rhs, $dot_dimension_numbers, $precision_config),
    (StableHLO_BroadcastInDimOp $bias, $dims)),
  (LiftAsTFXlaCallModule<"composite_dot_general_with_bias_fn">
    (ArgumentList $lhs, $rhs, $bias),
    (ResultList $res),
    (NamedAttributeList

// CHECK-NEXT: %[[conv:.*]] = "tf.Conv2D"(%arg0, %[[cst]])
// CHECK-NEXT: %[[bias:.*]] = "tf.AddV2"(%[[conv]], %[[cst_0]])
// CHECK-NEXT: return %[[bias]] : tensor<256x8x7x16xf32>
}

// CHECK-LABEL: convaddv2mul
func.func @convaddv2mul(%arg: tensor<256x32x32x3xf32>) -> tensor<256x8x7x16xf32> {
  %filter = arith.constant dense<2.0> : tensor<3x3x3x16xf32>
  %bias = arith.constant dense<3.0> : tensor<16xf32>

    filter_ = tf.random.uniform([2, 2, 1, 8])
    bias = tf.zeros([8])
    kwargs = {
        'input_': input_,
        'filter_': filter_,
        'bias': bias,
        'stride_w': 2,
        'stride_h': 2,
        'dilation_w': 1,
        'dilation_h': 1,
        'padding': 'SAME',
        'act': 'RELU'
    }

    self._assertOpAndComposite([input_, filter_, bias],

// Pattern rules for lifting ops with bias as functions
//===----------------------------------------------------------------------===//

def LiftDepthwiseConv2dNativeWithBias : Pat<
  (TF_BiasAddOp:$res
    (TF_DepthwiseConv2dNativeOp $input, $filter, $strides, $padding,
      $explicit_paddings, IsDataFormatNHWC:$data_format, $dilations),
    $bias, IsDataFormatNHWC:$bias_data_format),

    %1 = "tf.ReadVariableOp"(%0) {device = ""} : (tensor<!tf_type.resource<tensor<100x50xf32>>>) -> tensor<100x50xf32>
    %2 = "tf.VarHandleOp"() {_class = ["loc:@dense/bias"], allowed_devices = [], container = "", device = "", shared_name = "dense/bias"} : () -> tensor<!tf_type.resource<tensor<50xf32>>>
    %3 = "tf.ReadVariableOp"(%2) {device = ""} : (tensor<!tf_type.resource<tensor<50xf32>>>) -> tensor<50xf32>

  ),
  [(AreLastTwoDimsTransposed $perm_value), (IsNoneType $bias)]>;

// Fuses TFL_FullyConnectedOp and TFL_TransposeOp Rhs to TFL_BatchMatMulOp
def FuseTransposeFCRhsToBatchMatmul : Pat<
  (TFL_FullyConnectedOp
    2DTensorOf<[F32]>:$lhs,
    (TFL_TransposeOp TensorOf<[F32]>:$rhs, (Arith_ConstantOp:$perm_value $p0)),
    $bias, $TFL_AF_None, $TFL_FCWO_Default,
    $keep_num_dims, $asymmetric_quantize_inputs
    ),

			}
			t := k - bias
			if k <= bias {
				t = tmin
			} else if k >= bias+tmax {
				t = tmax
			}
			if digit < t {
				break
			}
			w, overflow = madd(0, w, base-t)
			if overflow {
				return "", punyError(encoded)
			}
		}
		if len(output) >= 1024 {
			return "", punyError(encoded)
		}
		x := int32(len(output) + 1)
		bias = adapt(i-oldI, x, oldI == 0)
		n += i / x

          $dilations, (location $conv)),
      [(CanFuseMulAndConv2D $filter_value, $mul_value), (HasOneUse $conv)]>;

// This rule does the following pattern match and rewrite:
//
//       input     bias                    input  value  bias  value
//          \      /                =>       \    /        \    /
//          BiasAdd    value                   Mul          Mul
//                \    /                           \       /