OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-06 12:20:52 +01:00
Code Issues Actions 38 Packages Projects Releases Wiki Activity

[pt2e][xnnpack_quantizer] quantize "mul" (#110428)

Browse Source 
Adding "mul" to list of partitions that are supported by the quantizer. This shows up in EDSR, where we still want to quantize the mul op

Differential Revision: [D49850151](https://our.internmc.facebook.com/intern/diff/D49850151/)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/110428
Approved by: https://github.com/jerryzh168
ghstack dependencies: #110427
This commit is contained in:
Max Ren 2023-10-03 15:06:00 -07:00 committed by PyTorch MergeBot
parent 66202ed29c
commit 08c7dcda65
2 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
torch/ao/quantization/pt2e/graph_utils.py
View File

@ -25,7 +25,7 @@ _EQUIVALENT_TYPES: List[Set] = [
{torch.nn.BatchNorm2d, torch.nn.functional.batch_norm}, {torch.nn.BatchNorm2d, torch.nn.functional.batch_norm},
{torch.nn.Hardtanh, torch.nn.functional.hardtanh, torch.nn.functional.hardtanh_}, {torch.nn.Hardtanh, torch.nn.functional.hardtanh, torch.nn.functional.hardtanh_},
{torch.add, operator.add, operator.iadd, "add", "add_"}, {torch.add, operator.add, operator.iadd, "add", "add_"},
{torch.mul, operator.mul, operator.imul}, {torch.mul, operator.mul, operator.imul, "mul", "mul_"},
] ]

2
torch/ao/quantization/quantizer/xnnpack_quantizer_utils.py
View File

@ -739,7 +739,7 @@ def _annotate_mul(
filter_fn: Optional[Callable[[Node], bool]] = None, filter_fn: Optional[Callable[[Node], bool]] = None,
) -> Optional[List[List[Node]]]: ) -> Optional[List[List[Node]]]:
mul_partitions = get_source_partitions( mul_partitions = get_source_partitions(
gm.graph, [operator.mul, torch.mul, operator.imul], filter_fn gm.graph, ["mul", "mul_", operator.mul, torch.mul, operator.imul], filter_fn
) )
mul_partitions = list(itertools.chain(*mul_partitions.values())) mul_partitions = list(itertools.chain(*mul_partitions.values()))
annotated_partitions = [] annotated_partitions = []