[Relax][PyTorch] Add rnn_tanh.input converter by cchung100m · Pull Request #19837 · apache/tvm

cchung100m · 2026-06-19T04:25:08Z

Hi Committers,

This PR addresses the rnn_tanh.input part of issue #18364. Any suggestions would be appreciated if you are available.

gemini-code-assist

Code Review

This pull request adds support for translating vanilla tanh-RNN cells (_rnn_tanh and _rnn_tanh_cell_unroll) in the TVM Relax PyTorch frontend. The review feedback identifies three critical issues: a typo in the function call premute_dims instead of permute_dims, an incorrect range definition for the reverse path iteration that lacks a step argument, and an incorrect initialization of the forward hidden state using zeros instead of slicing the input hidden state hx with take.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

cchung100m · 2026-06-21T10:57:41Z

/gemini review

gemini-code-assist

Code Review

This pull request adds support for translating PyTorch's vanilla tanh-RNN (rnn_tanh.input) in the Relax frontend, including support for bidirectional and batch-first configurations, along with corresponding unit tests. The review feedback suggests optimizing the unrolled RNN cell loop by pre-computing the sum of the input and hidden biases outside the loop to avoid redundant additions.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

tlopex

Thanks @cchung100m The added test exercises the decomposed graph, not the new rnn_tanh.input converter. With run_ep_decomposition=False, _rnn_tanh returns only output, so getitem(..., 1) incorrectly returns the output sequence instead of h_n. Please return (output, h_n) and add a no-decomposition test covering the hidden-state output.

cchung100m · 2026-06-22T14:06:08Z

/gemini review

gemini-code-assist

Code Review

This pull request adds support for translating PyTorch's rnn_tanh operator in the TVM Relax frontend. It implements the _rnn_tanh translation logic and unrolls the RNN cells using _rnn_tanh_cell_unroll for both unidirectional and bidirectional configurations. The review feedback highlights a critical bug where relax.op.Tuple is used instead of relax.Tuple, recommends raising a NotImplementedError for dynamic sequence lengths since the unrolling loop requires an integer sequence length, and suggests adding a decorator to skip the new test if LLVM is not enabled.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

cchung100m · 2026-06-27T15:09:45Z

/gemini review

gemini-code-assist

Code Review

This pull request adds support for vanilla tanh-RNN cells in the PyTorch exported program translator by implementing the _rnn_tanh_cell_unroll and _rnn_tanh methods, and registers the new operator in the conversion map. A comprehensive unit test is also added to verify different configurations. The feedback points out that retrieving shapes and data types via raw FX node metadata is fragile and inconsistent with other translator implementations; it is recommended to query the translated Relax expressions directly using self.shape_of and struct_info.dtype instead.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

gemini-code-assist · 2026-06-27T15:11:01Z

+        def _node_meta(fx_node):
+            meta = fx_node.meta
+            return meta["val"] if "val" in meta else meta["tensor_meta"]
+
+        input_meta = _node_meta(node.args[0])
+        input_shape = list(input_meta.shape)
+        if batch_first:
+            batch_size, seq_len, input_size = input_shape
+        else:
+            seq_len, batch_size, input_size = input_shape
+
+        if not isinstance(seq_len, int):
+            raise NotImplementedError("Dynamic sequence length is not supported for rnn_tanh")
+
+        # params per direction: weight_ih, weight_hh, [bias_ih, bias_hh]
+        params_per_direction = 4 if has_biases else 2
+
+        # A vanilla RNN has a single gate, so weight_ih has shape (hidden_size, input_size)
+        if params and len(params) >= 2:
+            hidden_size = int(_node_meta(node.args[2][0]).shape[0])
+        else:
+            hidden_size = 16
+
+        dtype = self._convert_data_type(input_meta.dtype)


The _node_meta helper function accesses the raw FX node's metadata (node.args[0].meta) to retrieve shapes and dtypes. This is fragile and inconsistent with the rest of the translator (such as _lstm and _gru), which safely query the translated Relax expressions using self.shape_of(input_tensor) and input_tensor.struct_info.dtype. Querying the Relax expressions directly is much more robust and allows us to completely remove the _node_meta helper.

input_shape = self.shape_of(input_tensor) if batch_first: batch_size, seq_len, input_size = input_shape else: seq_len, batch_size, input_size = input_shape seq_len = int(seq_len) if isinstance(seq_len, tvm.tirx.IntImm) else seq_len batch_size = int(batch_size) if isinstance(batch_size, tvm.tirx.IntImm) else batch_size input_size = int(input_size) if isinstance(input_size, tvm.tirx.IntImm) else input_size if not isinstance(seq_len, int): raise NotImplementedError("Dynamic sequence length is not supported for rnn_tanh") # params per direction: weight_ih, weight_hh, [bias_ih, bias_hh] params_per_direction = 4 if has_biases else 2 # A vanilla RNN has a single gate, so weight_ih has shape (hidden_size, input_size) if params and len(params) >= 2: hidden_size = int(self.shape_of(params[0])[0]) else: hidden_size = 16 dtype = input_tensor.struct_info.dtype

Thanks for the suggestion. I looked into this and kept _node_meta intentionally - reverting to self_shape_of(input_tensor) / input_tensor.struct_info.dtype would re-break the now-passing test_rnn_tanh.

[2026-06-22T16:46:09.119Z] =================================== FAILURES ===================================
[2026-06-22T16:46:09.119Z] ________________________________ test_rnn_tanh _________________________________
[2026-06-22T16:46:09.119Z] [gw0] linux -- Python 3.10.19 /venv/apache-tvm-py3.10/bin/python3
[2026-06-22T16:46:09.119Z] tests/python/relax/test_frontend_from_exported_program.py:8604: in test_rnn_tanh
[2026-06-22T16:46:09.119Z] _check(
[2026-06-22T16:46:09.119Z] tests/python/relax/test_frontend_from_exported_program.py:8587: in _check
[2026-06-22T16:46:09.119Z] mod = from_exported_program(exported_program, run_ep_decomposition=False)
[2026-06-22T16:46:09.119Z] python/tvm/relax/frontend/torch/exported_program_translator.py:2389: in from_exported_program
[2026-06-22T16:46:09.119Z] return ExportedProgramImporter().from_exported_program(
[2026-06-22T16:46:09.119Z] python/tvm/relax/frontend/torch/exported_program_translator.py:2224: in from_exported_program
[2026-06-22T16:46:09.119Z] output_args = self._translate_fx_graph(
[2026-06-22T16:46:09.119Z] python/tvm/relax/frontend/torch/exported_program_translator.py:1533: in _translate_fx_graph
[2026-06-22T16:46:09.119Z] self.env[node] = self.convert_mapfunc_name
[2026-06-22T16:46:09.119Z] python/tvm/relax/frontend/torch/exported_program_translator.py:1007: in _rnn_tanh
[2026-06-22T16:46:09.119Z] dtype = input_tensor.struct_info.dtype
[2026-06-22T16:46:09.119Z] E AttributeError: 'Var' object has no attribute 'struct_info'
[2026-06-22T16:46:09.119Z] =============================== warnings summary ===============================

cchung100m · 2026-06-27T17:51:44Z

Hi @tlopex

Thanks for your friendly reminder, I updated the part you mentioned.

tlopex

LGTM

cchung100m · 2026-06-28T03:19:26Z

Thanks to @tlopex

gemini-code-assist Bot reviewed Jun 19, 2026

View reviewed changes

Comment thread python/tvm/relax/frontend/torch/exported_program_translator.py Outdated

Comment thread python/tvm/relax/frontend/torch/exported_program_translator.py Outdated

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

[Relax][PyTorch] Add rnn_tanh.input converter

4cebdbd

cchung100m force-pushed the issue-18364-rnn_tanh branch from 78f1250 to 4cebdbd Compare June 19, 2026 04:28

[Relax][PyTorch] Add test case: test_rnn_tanh

78ace25

cchung100m force-pushed the issue-18364-rnn_tanh branch from d6bfdc3 to 78ace25 Compare June 21, 2026 10:24

gemini-code-assist Bot reviewed Jun 21, 2026

View reviewed changes

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

cchung100m marked this pull request as ready for review June 21, 2026 14:54

tlopex requested changes Jun 22, 2026

View reviewed changes

[Relax][PyTorch] Return (output, h_n) instead of only output

548699e

cchung100m marked this pull request as draft June 22, 2026 12:40

[Relax][PyTorch] Bias folding

dad5a74

gemini-code-assist Bot reviewed Jun 22, 2026

View reviewed changes

Comment thread python/tvm/relax/frontend/torch/exported_program_translator.py Outdated

Comment thread python/tvm/relax/frontend/torch/exported_program_translator.py Outdated

Comment thread tests/python/relax/test_frontend_from_exported_program.py

[Relax][PyTorch] Fix relax.op.Tuple

48d4564

cchung100m force-pushed the issue-18364-rnn_tanh branch 4 times, most recently from ad6eb44 to 4fdb43c Compare June 27, 2026 11:10

Add the metadata introspection

14387b5

cchung100m force-pushed the issue-18364-rnn_tanh branch from 4fdb43c to 14387b5 Compare June 27, 2026 14:12

gemini-code-assist Bot reviewed Jun 27, 2026

View reviewed changes

cchung100m marked this pull request as ready for review June 27, 2026 15:45

cchung100m requested a review from tlopex June 27, 2026 15:45

tlopex approved these changes Jun 28, 2026

View reviewed changes

tlopex merged commit 182db0f into apache:main Jun 28, 2026
13 checks passed

cchung100m deleted the issue-18364-rnn_tanh branch June 28, 2026 03:19

Uh oh!

Conversation

cchung100m commented Jun 19, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

cchung100m commented Jun 21, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

tlopex left a comment

Choose a reason for hiding this comment

Uh oh!

cchung100m commented Jun 22, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

cchung100m commented Jun 27, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Jun 27, 2026

Choose a reason for hiding this comment

Uh oh!

cchung100m Jun 27, 2026

Choose a reason for hiding this comment

Uh oh!

cchung100m commented Jun 27, 2026

Uh oh!

tlopex left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

cchung100m commented Jun 28, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants