Skip to content

Commit a8c7580

Browse files
authored
[Relax][PyTorch] Enable run_ep_decomposition by default (#18471)
## Why We have finished the migration for our tests then we could set default to run ep decompose. ## How Update tests and exported_program_translator.py
1 parent b2c58ef commit a8c7580

File tree

2 files changed

+259
-280
lines changed

2 files changed

+259
-280
lines changed

python/tvm/relax/frontend/torch/exported_program_translator.py

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -208,6 +208,15 @@ def _native_group_norm(self, node: fx.Node) -> relax.Var:
208208
)
209209
)
210210

211+
def _native_layer_norm(self, node: fx.Node) -> relax.Var:
212+
# native_layer_norm signature: (input, normalized_shape, weight, bias, eps)
213+
x = self.env[node.args[0]]
214+
normalized_shape = node.args[1]
215+
gamma = self.env.get(node.args[2], None) if len(node.args) > 2 else None
216+
beta = self.env.get(node.args[3], None) if len(node.args) > 3 else None
217+
eps = node.args[4] if len(node.args) > 4 else 1e-05
218+
return self._layer_norm_impl(x, gamma, beta, eps, normalized_shape)
219+
211220
def _upsample_impl(
212221
self,
213222
x: relax.Expr,
@@ -1058,6 +1067,7 @@ def create_convert_map(
10581067
"instance_norm.default": self._instance_norm,
10591068
"native_group_norm.default": self._native_group_norm,
10601069
"layer_norm.default": self._layer_norm,
1070+
"native_layer_norm.default": self._native_layer_norm,
10611071
"linear.default": self._linear,
10621072
"lstm.input": self._lstm,
10631073
"gru.input": self._gru,
@@ -1403,7 +1413,7 @@ def from_exported_program(
14031413
keep_params_as_input: bool = False,
14041414
unwrap_unit_return_tuple: bool = False,
14051415
no_bind_return_tuple: bool = False,
1406-
run_ep_decomposition: bool = False,
1416+
run_ep_decomposition: bool = True,
14071417
) -> tvm.IRModule:
14081418
"""Convert a PyTorch ExportedProgram to a Relax program
14091419
@@ -1426,8 +1436,7 @@ def from_exported_program(
14261436
run_ep_decomposition : bool
14271437
A boolean flag indicating whether to run PyTorch's decomposition on the
14281438
exported program before translation. When True, high-level operators will
1429-
be decomposed into their constituent parts. Defaults to False for backward
1430-
compatibility.
1439+
be decomposed into their constituent parts. Defaults to True.
14311440
14321441
Returns
14331442
-------

0 commit comments

Comments
 (0)