Skip to content

Commit 21095b9

Browse files
authored
[Relax][PyTorch] Add support for abs, ceil, erf, floor, log ops and refactor unary tests (#17622)
* Update fx_translator.py * Update test_frontend_from_fx.py * Update test_frontend_from_fx.py * Update test_frontend_from_fx.py * Update test_frontend_from_fx.py * Update test_frontend_from_fx.py * Update erf ops * Update template for unary ops * Update test_frontend_from_fx.py * delete test_name * Update test_frontend_from_fx.py * Update test_frontend_from_fx.py * Update test_frontend_from_fx.py
1 parent ce22736 commit 21095b9

File tree

2 files changed

+239
-561
lines changed

2 files changed

+239
-561
lines changed

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -616,21 +616,26 @@ def create_convert_map(
616616
nn.Flatten: self._flatten_module,
617617
## call_function and call_method
618618
# unary
619+
"abs": self._unary_op(relax.op.abs),
619620
"acos": self._unary_op(relax.op.acos),
620621
"acosh": self._unary_op(relax.op.acosh),
621622
"asin": self._unary_op(relax.op.asin),
622623
"asinh": self._unary_op(relax.op.asinh),
623624
"atan": self._unary_op(relax.op.atan),
624625
"atanh": self._unary_op(relax.op.atanh),
626+
"ceil": self._unary_op(relax.op.ceil),
625627
"clamp": self._clamp,
626628
"cos": self._unary_op(relax.op.cos),
627629
"cosh": self._unary_op(relax.op.cosh),
628630
"dropout": lambda node: self.env[node.args[0]],
631+
"erf": self._unary_op(relax.op.erf),
629632
"exp": self._unary_op(relax.op.exp),
633+
"floor": self._unary_op(relax.op.floor),
630634
"gelu": self._gelu,
631635
"hardsigmoid": self._hardsigmoid,
632636
"hardswish": self._hardswish,
633637
"leaky_relu": self._leakyrelu,
638+
"log": self._unary_op(relax.op.log),
634639
"log_softmax": self._log_softmax,
635640
"neg": self._unary_op(relax.op.negative),
636641
"relu": self._unary_op(relax.op.nn.relu),

0 commit comments

Comments
 (0)