@@ -120,7 +120,7 @@ class SharedSpecPattern(QuantizationPattern):
120120 """
121121
122122 @abstractmethod
123- def partition_types (self ) -> list [torch . nn . Module ]:
123+ def partition_types (self ) -> list [OpOverload ]:
124124 pass
125125
126126 @staticmethod
@@ -152,6 +152,52 @@ def get_anchors(
152152 return self .get_shared_spec_anchors (gm , fused_partition )
153153
154154
155+ class SharedSpecMultipleInputPattern (QuantizationPattern ):
156+ """
157+ Quantization pattern for shared quantization with multiple inputs.
158+
159+ The quantization is derived from the previous node quantization and inputs and the output shares the same
160+ quantization parameters (scale and zero-point).
161+ """
162+
163+ @abstractmethod
164+ def partition_types (self ) -> list [OpOverload ]:
165+ pass
166+
167+ @staticmethod
168+ def get_shared_spec_anchors (
169+ gm : fx .GraphModule , fused_partition : list [fx .GraphModule ]
170+ ) -> PartitionAnchors | None :
171+ node = fused_partition [0 ].nodes [- 1 ]
172+
173+ quantized_input = None
174+ for prev_node in node .args :
175+ if Q_ANNOTATION_KEY in prev_node .meta :
176+ quantized_input = prev_node
177+ break
178+
179+ if quantized_input is not None :
180+ inputs = []
181+ for idx , _ in enumerate (node .args ):
182+ inputs .append (
183+ (node , NodeArgsIdx (idx ), SharedQuantizationSpec (quantized_input ))
184+ )
185+ else :
186+ return None
187+
188+ return PartitionAnchors (
189+ inputs = inputs ,
190+ weights = [],
191+ biases = [],
192+ output = [(node , SharedQuantizationSpec (quantized_input ))],
193+ )
194+
195+ def get_anchors (
196+ self , gm : fx .GraphModule , fused_partition : list [fx .GraphModule ]
197+ ) -> PartitionAnchors | None :
198+ return self .get_shared_spec_anchors (gm , fused_partition )
199+
200+
155201class SingleInputBasicPattern (QuantizationPattern ):
156202 @abstractmethod
157203 def partition_types (self ) -> list [OpOverload ]:
@@ -300,7 +346,7 @@ class AddTensorPattern(QuantizationPattern):
300346 Basic quantization for all inputs and output.
301347 """
302348
303- def partition_types (self ) -> list [torch . nn . Module ]:
349+ def partition_types (self ) -> list [OpOverload ]:
304350 return [torch .ops .aten .add .Tensor ]
305351
306352 def get_anchors (
@@ -333,7 +379,7 @@ class BMMPattern(QuantizationPattern):
333379 Quantizer for BatchMatMul operator.
334380 """
335381
336- def partition_types (self ) -> list [torch . nn . Module ]:
382+ def partition_types (self ) -> list [OpOverload ]:
337383 return [torch .ops .aten .bmm .default ]
338384
339385 def get_anchors (
@@ -372,7 +418,7 @@ class SubTensorPattern(QuantizationPattern):
372418 Basic quantization for all inputs and output.
373419 """
374420
375- def partition_types (self ) -> list [torch . nn . Module ]:
421+ def partition_types (self ) -> list [OpOverload ]:
376422 return [torch .ops .aten .sub .Tensor ]
377423
378424 def get_anchors (
@@ -426,7 +472,7 @@ def get_anchors(
426472
427473 quantized_input = None
428474 for prev_node in node .args [0 ]:
429- if "quantization_annotation" in prev_node .meta :
475+ if Q_ANNOTATION_KEY in prev_node .meta :
430476 quantized_input = prev_node
431477 break
432478
@@ -876,6 +922,17 @@ def partition_types(self):
876922 return [torch .ops .aten .log .default ]
877923
878924
925+ class MaximumPattern (SharedSpecMultipleInputPattern ):
926+ """
927+ Quantization pattern for Maximum quantization. Accepts 1 or 2 input nodes.
928+
929+ Basic quantization for all inputs and output.
930+ """
931+
932+ def partition_types (self ) -> list [OpOverload ]:
933+ return [torch .ops .aten .maximum .default ]
934+
935+
879936class MaxPool1DPattern (SharedSpecPattern ):
880937 """Quantizer for the MaxPool1D operator."""
881938
0 commit comments