Commit 0b027d2
Fix DecomposeLayerNormPass to handle 6-arg layer_norm
Summary:
## Problem
When using `nn.LayerNorm` in models that go through the ARM backend's quantization flow, the `DecomposeLayerNormPass` fails with:
```
ValueError: DecomposeLayerNormPass: too many values to unpack (expected 2)
```
This happens because `torch.ops.aten.layer_norm.default` has **6 arguments**:
```
layer_norm(input, normalized_shape, weight, bias, eps, cudnn_enable)
```
But `DecomposeLayerNormPass` only handled up to 5 arguments (for `native_layer_norm`).
The error occurs during `transform_for_annotation_pipeline` in the ARM quantizer, which runs before edge transformation when the op is still `aten.layer_norm.default`.
## Solution
Add `case 6:` to the `match len(args)` block in `DecomposeLayerNormPass.call()` to handle the 6th argument (`cudnn_enable`). This argument is simply ignored during decomposition since it's only relevant for cuDNN GPU optimization.
---
> Generated by [Confucius Code Assist (CCA)](https://www.internalfb.com/wiki/Confucius/Analect/Shared_Analects/Confucius_Code_Assist_(CCA)/)
[Confucius Session](https://www.internalfb.com/confucius?host=92481.od.fbinfra.net&port=8086&tab=Chat&session_id=eace3d92-ed78-11f0-b67c-c7843469b0d5&entry_name=Code+Assist), [Trace](https://www.internalfb.com/confucius?session_id=eace3d92-ed78-11f0-b67c-c7843469b0d5&tab=Trace)
Differential Revision: D903957861 parent 5b4900c commit 0b027d2
File tree
2 files changed
+6
-1
lines changed- backends/arm
- _passes
- test
2 files changed
+6
-1
lines changed| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
90 | 90 | | |
91 | 91 | | |
92 | 92 | | |
| 93 | + | |
| 94 | + | |
| 95 | + | |
| 96 | + | |
| 97 | + | |
93 | 98 | | |
94 | 99 | | |
95 | 100 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
19 | 19 | | |
20 | 20 | | |
21 | 21 | | |
22 | | - | |
| 22 | + | |
23 | 23 | | |
24 | 24 | | |
25 | 25 | | |
| |||
0 commit comments