This is an automated email from the ASF dual-hosted git repository.
syfeng pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/unity by this push:
new ec1184e105 [Unity] [Bugfix] Fix KeyError:'None' in layer_norm and
correctly use the normalized_shape (#15894)
ec1184e105 is described below
commit ec1184e1059fc7fcbe970614f3e9dcc4d4687304
Author: Thrsu <[email protected]>
AuthorDate: Mon Oct 9 22:36:15 2023 +0800
[Unity] [Bugfix] Fix KeyError:'None' in layer_norm and correctly use the
normalized_shape (#15894)
* Fix the KeyError and correctly use the normalized_shape
* Update test_frontend_from_fx.py
---
python/tvm/relax/frontend/torch/fx_translator.py | 9 ++++--
tests/python/relax/test_frontend_from_fx.py | 36 ++++++++++++++++++++++++
2 files changed, 43 insertions(+), 2 deletions(-)
diff --git a/python/tvm/relax/frontend/torch/fx_translator.py
b/python/tvm/relax/frontend/torch/fx_translator.py
index d08e8858dc..a52b73f3ab 100644
--- a/python/tvm/relax/frontend/torch/fx_translator.py
+++ b/python/tvm/relax/frontend/torch/fx_translator.py
@@ -981,10 +981,15 @@ class TorchFXImporter:
dim_num = len(normalized_shape)
axes = list(range(-dim_num, 0))
- gamma = self.env[node.kwargs["weight"]]
+ gamma = node.kwargs["weight"]
+ if gamma is None:
+ shape_tuple = [int(s) for s in normalized_shape]
+ gamma = relax.const(np.ones(shape_tuple), x.struct_info.dtype)
+ else:
+ gamma = self.env[gamma]
beta = node.kwargs["bias"]
if beta is None:
- shape_tuple = [int(s) for s in normalized_shape.values]
+ shape_tuple = [int(s) for s in normalized_shape]
beta = relax.const(np.zeros(shape_tuple), x.struct_info.dtype)
else:
beta = self.env[beta]
diff --git a/tests/python/relax/test_frontend_from_fx.py
b/tests/python/relax/test_frontend_from_fx.py
index abe6c947e4..c815051076 100644
--- a/tests/python/relax/test_frontend_from_fx.py
+++ b/tests/python/relax/test_frontend_from_fx.py
@@ -1073,6 +1073,8 @@ def test_layernorm():
def test_functional_layernorm():
+ import numpy as np
+
input_info = [([1, 3, 10, 10], "float32")]
class LayerNorm(Module):
@@ -1116,6 +1118,40 @@ def test_functional_layernorm():
}
verify_model(model, input_info, binding, expected1)
+ class LayerNorm2(Module):
+ def __init__(self, shape):
+ super().__init__()
+ self.shape = shape
+ self.weight = None
+ self.bias = None
+
+ def forward(self, input):
+ return torch.nn.functional.layer_norm(input, self.shape,
self.weight, self.bias, 1e-5)
+
+ @tvm.script.ir_module
+ class expected2:
+ @R.function
+ def main(
+ input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
+ ) -> R.Tensor((1, 3, 10, 10), dtype="float32"):
+ with R.dataflow():
+ lv: R.Tensor((1, 3, 10, 10), dtype="float32") =
R.nn.layer_norm(
+ input_1,
+ gamma=relax.const(np.ones((10, 10)), dtype="float32"),
+ beta=relax.const(np.zeros((10, 10)), dtype="float32"),
+ axes=[-2, -1],
+ epsilon=1e-05,
+ center=True,
+ scale=True,
+ )
+ gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
+ R.output(gv)
+ return gv
+
+ model = LayerNorm2((10, 10))
+ binding = {}
+ verify_model(model, input_info, binding, expected2)
+
def test_cross_entropy():
input_info = [([3, 2], "float32"), ([3], "int32")]
