This is an automated email from the ASF dual-hosted git repository.
syfeng pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new e74ded2163 [Relax][PyTorch] Add support for log2, log10 and log1p ops
(#17766)
e74ded2163 is described below
commit e74ded2163336f1c41ddc1767c659b62ce4b6de8
Author: Shushi Hong <[email protected]>
AuthorDate: Thu Mar 20 13:09:48 2025 +0800
[Relax][PyTorch] Add support for log2, log10 and log1p ops (#17766)
* Update fx_translator.py
* Update test_frontend_from_fx.py
* Update fx_translator.py
* Update test_frontend_from_fx.py
* Update test_frontend_from_fx.py
---
python/tvm/relax/frontend/torch/fx_translator.py | 20 +++++++
tests/python/relax/test_frontend_from_fx.py | 71 +++++++++++++++++++++++-
2 files changed, 88 insertions(+), 3 deletions(-)
diff --git a/python/tvm/relax/frontend/torch/fx_translator.py
b/python/tvm/relax/frontend/torch/fx_translator.py
index 9b835b0eee..98de2e114b 100644
--- a/python/tvm/relax/frontend/torch/fx_translator.py
+++ b/python/tvm/relax/frontend/torch/fx_translator.py
@@ -68,6 +68,23 @@ class TorchFXImporter(BaseFXGraphImporter):
alpha = module.negative_slope
return self.block_builder.emit(relax.op.nn.leakyrelu(x, alpha))
+ def _log2(self, node: fx.Node) -> relax.Var:
+ x = self.env[node.args[0]]
+ return self.block_builder.emit(
+ relax.op.divide(relax.op.log(x), relax.const(0.6931471805599453,
x.struct_info.dtype))
+ )
+
+ def _log10(self, node: fx.Node) -> relax.Var:
+ x = self.env[node.args[0]]
+ return self.block_builder.emit(
+ relax.op.divide(relax.op.log(x), relax.const(2.302585092994046,
x.struct_info.dtype))
+ )
+
+ def _log1p(self, node: fx.Node) -> relax.Var:
+ x = self.env[node.args[0]]
+ one = relax.const(1, x.struct_info.dtype)
+ return self.block_builder.emit(relax.op.log(relax.op.add(x, one)))
+
def _log_softmax_module(self, node: fx.Node) -> relax.Var:
x = self.env[node.args[0]]
module = self.named_modules[node.target]
@@ -685,6 +702,9 @@ class TorchFXImporter(BaseFXGraphImporter):
"isnan": self._unary_op(relax.op.isnan),
"leaky_relu": self._leakyrelu,
"log": self._unary_op(relax.op.log),
+ "log2": self._log2,
+ "log10": self._log10,
+ "log1p": self._log1p,
"logical_not": self._unary_op(relax.op.logical_not),
"log_softmax": self._log_softmax,
"neg": self._unary_op(relax.op.negative),
diff --git a/tests/python/relax/test_frontend_from_fx.py
b/tests/python/relax/test_frontend_from_fx.py
index f06ce7a753..962a6accff 100644
--- a/tests/python/relax/test_frontend_from_fx.py
+++ b/tests/python/relax/test_frontend_from_fx.py
@@ -2128,9 +2128,6 @@ def test_extended_unary_ops():
verify_model(Gelu(), input_info, {}, expected_gelu)
verify_model(Gelu2(), input_info, {}, expected_gelu)
- # leaky_relu
- test_leakyrelu()
-
# hardsigmoid
class Hardsigmoid(torch.nn.Module):
def __init__(self):
@@ -2226,6 +2223,74 @@ def test_extended_unary_ops():
verify_model(Hardtanh(), input_info, {}, expected1)
verify_model(Hardtanh2(), input_info, {}, expected1)
+ # leaky_relu
+ test_leakyrelu()
+
+ # log2
+ class Log2(Module):
+ def forward(self, x):
+ return torch.log2(x)
+
+ @tvm.script.ir_module
+ class Expected_log2:
+ @R.function
+ def main(
+ inp_0: 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.log(inp_0)
+ lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.divide(
+ lv, R.const(0.6931471805599453, "float32")
+ )
+ gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv1
+ R.output(gv)
+ return gv
+
+ verify_model(Log2(), input_info, {}, Expected_log2)
+
+ # log10
+ class Log10(Module):
+ def forward(self, x):
+ return torch.log10(x)
+
+ @tvm.script.ir_module
+ class Expected_log10:
+ @R.function
+ def main(
+ inp_0: 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.log(inp_0)
+ lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.divide(
+ lv, R.const(2.302585092994046, "float32")
+ )
+ gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv1
+ R.output(gv)
+ return gv
+
+ verify_model(Log10(), input_info, {}, Expected_log10)
+
+ # log1p
+ class Log1p(Module):
+ def forward(self, x):
+ return torch.log1p(x)
+
+ @tvm.script.ir_module
+ class Expected_log1p:
+ @R.function
+ def main(
+ inp_0: 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.log(
+ R.add(inp_0, R.const(1, "float32"))
+ )
+ gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
+ R.output(gv)
+ return gv
+
+ verify_model(Log1p(), input_info, {}, Expected_log1p)
+
# logical_not
class LogicalNot(Module):
def forward(self, input):
