Quantizer for BatchMatMul operator.
| 303 | |
| 304 | |
| 305 | class BMMPattern(QuantizationPattern): |
| 306 | """ |
| 307 | Quantizer for BatchMatMul operator. |
| 308 | """ |
| 309 | |
| 310 | def partition_types(self) -> list[torch.nn.Module]: |
| 311 | return [torch.ops.aten.bmm.default] |
| 312 | |
| 313 | def get_anchors( |
| 314 | self, gm: fx.GraphModule, fused_partition: list[fx.GraphModule] |
| 315 | ) -> PartitionAnchors | None: |
| 316 | bmm_node = fused_partition[0].nodes[-1] |
| 317 | |
| 318 | return PartitionAnchors( |
| 319 | inputs=[ |
| 320 | (bmm_node, NodeArgsIdx(0)), |
| 321 | (bmm_node, NodeArgsIdx(1)), |
| 322 | ], |
| 323 | biases=[], |
| 324 | output=[(bmm_node,)], |
| 325 | ) |
| 326 | |
| 327 | |
| 328 | class SubTensorPattern(QuantizationPattern): |