This is an automated email from the ASF dual-hosted git repository.
tlopex 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 9f8cead697 [Relax][PyTorch] Support aten.scatter.src in
ExportedProgram importer (#19935)
9f8cead697 is described below
commit 9f8cead6972eaeb22e3663e3108144b225b8260c
Author: Hangshuai He <[email protected]>
AuthorDate: Sat Sep 12 13:41:39 2026 +0800
[Relax][PyTorch] Support aten.scatter.src in ExportedProgram importer
(#19935)
This PR adds PyTorch ExportedProgram frontend support for
aten.scatter.src by mapping it to the existing scatter conversion path.
The importer already supports scatter.value, and BaseFXGraphImporter
already provides _scatter for tensor-source scatter. This patch
registers scatter.src in the ExportedProgram convert map and adds a
regression test.
Test:
python -m pytest
tests/python/relax/test_frontend_from_exported_program.py -k "scatter"
-q
---
.../frontend/torch/exported_program_translator.py | 1 +
.../relax/test_frontend_from_exported_program.py | 27 ++++++++++++++++++++++
2 files changed, 28 insertions(+)
diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py
b/python/tvm/relax/frontend/torch/exported_program_translator.py
index f01831d63c..2029dcc121 100644
--- a/python/tvm/relax/frontend/torch/exported_program_translator.py
+++ b/python/tvm/relax/frontend/torch/exported_program_translator.py
@@ -1774,6 +1774,7 @@ class ExportedProgramImporter(BaseFXGraphImporter):
"round.decimals": self._round,
"rsqrt.default": self._rsqrt,
"scalar_tensor.default": self._scalar_tensor,
+ "scatter.src": self._scatter,
"scatter.value": self._scatter_value,
"rsub.Tensor": self._rsub,
"rsub.Scalar": self._rsub,
diff --git a/tests/python/relax/test_frontend_from_exported_program.py
b/tests/python/relax/test_frontend_from_exported_program.py
index 565fd836bd..c74573c49f 100644
--- a/tests/python/relax/test_frontend_from_exported_program.py
+++ b/tests/python/relax/test_frontend_from_exported_program.py
@@ -9083,6 +9083,33 @@ def test_scatter_value():
verify_model(ScatterValue(), example_args, {}, Expected)
+def test_scatter_src():
+ class ScatterSrc(Module):
+ def forward(self, x, index, src):
+ return x.scatter(1, index, src)
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(
+ x: R.Tensor((4, 8), dtype="float32"),
+ index: R.Tensor((4, 2), dtype="int64"),
+ src: R.Tensor((4, 2), dtype="float32"),
+ ) -> R.Tuple(R.Tensor((4, 8), dtype="float32")):
+ with R.dataflow():
+ lv: R.Tensor((4, 8), dtype="float32") = R.scatter_elements(x,
index, src, axis=1)
+ gv: R.Tuple(R.Tensor((4, 8), dtype="float32")) = (lv,)
+ R.output(gv)
+ return gv
+
+ example_args = (
+ torch.randn(4, 8, dtype=torch.float32),
+ torch.randint(0, 8, (4, 2), dtype=torch.int64),
+ torch.randn(4, 2, dtype=torch.float32),
+ )
+ verify_model(ScatterSrc(), example_args, {}, Expected)
+
+
def test_grid_sample():
class GridSample(Module):
def forward(self, input, grid):
