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>

// 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),

  ),
  [(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
    ),

        auto attr = rewriter.getZeroAttr(type);
        bias = rewriter.create<arith::ConstantOp>(add_op.getLoc(), type, attr);
        auto none_af = rewriter.getStringAttr("NONE");

        bias =
            rewriter.create<AddOp>(add_op.getLoc(), bias, constant_val, none_af)
                .getOutput();
      } else {
        // If there no pre-existing bias and the `constant_val` is 1D, simply

              when {
                k <= bias -> TMIN
                k >= bias + TMAX -> TMAX
                else -> k - bias
              }
            if (q < t) break
            result.writeByte((t + ((q - t) % (BASE - t))).punycodeDigit)
            q = (q - t) / (BASE - t)
          }

          result.writeByte(q.punycodeDigit)
          bias = adapt(delta, h + 1, h == b)
          delta = 0

    (TF_BiasAddOp:$bias_add
      $conv_out,
      (TF_ConstOp:$bias IsFloatElementsAttr:$bias_value), $data_format),
    (TF_ConstOp:$add_rhs IsFloatElementsAttr:$add_rhs_value)),
  (TF_BiasAddOp
    $conv_out, (TF_AddV2Op $bias, (ReshapeTo1DTensor $add_rhs)), $data_format),
  [(HasOneUse $bias_add),
   (ReshapableTo1DTensor $add_rhs),
   (HasEqualElementSize<[-1], [-1]> $bias, $add_rhs)]>;

// Fuse AffineOp followed by an MulOp patterns.

      : public QuantizationSpecTraitBase<
            ConcreteType, AccumulatorUniformScale<Bias, Operands...>::Impl> {
   public:
    // Whether the index-th operand is a bias.
    static bool IsBias(int index) { return index == Bias; }

    // Returns the indexes of all the non-bias operands.
    static std::vector<int> GetAllNonBiasOperands() {
      return std::vector<int>({Operands...});
    }
  };
};