Skip to content

[Refactor] Simplify LSTM/GRUModule recurrent-mode shape normalization #3768

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
vmoens merged 4 commits into from
May 18, 2026
Merged

[Refactor] Simplify LSTM/GRUModule recurrent-mode shape normalization #3768

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
e69f12c
Update
vmoens May 15, 2026
6a48381
Update
vmoens May 17, 2026
f62f782
Update
vmoens May 18, 2026
ad40fcd
Update
vmoens May 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions test/modules/test_rnn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -662,6 +662,10 @@ def test_lstm_vmap_complex_model(self):
in_key="embed",
out_key="features",
python_based=True,
# vmap cannot trace through ``torch._higher_order_ops.scan``; the
# 'pad' backend keeps the time loop as a plain Python call into
# the Python-based LSTM, which is fully vmap-compatible.
recurrent_backend="pad",
)
mlp = TensorDictModule(
MLP(
Expand Down Expand Up @@ -2004,6 +2008,10 @@ def test_gru_vmap_complex_model(self):
in_key="embed",
out_key="features",
python_based=True,
# vmap cannot trace through ``torch._higher_order_ops.scan``; the
# 'pad' backend keeps the time loop as a plain Python call into
# the Python-based GRU, which is fully vmap-compatible.
recurrent_backend="pad",
)
mlp = TensorDictModule(
MLP(
Expand Down
78 changes: 37 additions & 41 deletions torchrl/modules/tensordict_module/rnn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
from tensordict import TensorDict, TensorDictBase, unravel_key_list
from tensordict.base import NO_DEFAULT
from tensordict.nn import dispatch, TensorDictModuleBase as ModuleBase
from tensordict.utils import expand_as_right, prod, set_lazy_legacy
from tensordict.utils import expand_as_right, set_lazy_legacy
from torch import nn, Tensor
from torch.nn.modules.rnn import RNNCellBase

Expand Down Expand Up @@ -915,28 +915,34 @@ def forward(self, tensordict: TensorDictBase):
# we want to get an error if the value input is missing, but not the hidden states
defaults = [NO_DEFAULT, None, None]
shape = tensordict.shape
tensordict_shaped = tensordict
if self.recurrent_mode:
# if less than 2 dims, unsqueeze
ndim = tensordict_shaped.get(self.in_keys[0]).ndim
while ndim < 3:
tensordict_shaped = tensordict_shaped.unsqueeze(0)
ndim += 1
if ndim > 3:
dims_to_flatten = ndim - 3
# we assume that the tensordict can be flattened like this
nelts = prod(tensordict_shaped.shape[: dims_to_flatten + 1])
tensordict_shaped = tensordict_shaped.apply(
lambda value: value.flatten(0, dims_to_flatten),
batch_size=[nelts, tensordict_shaped.shape[-1]],
# Straight-line shape normalization. Time is the last batch dim;
# all earlier batch dims are folded into a single leading B.
# Cheaper and simpler than the historical ``while ndim < 3`` loop
# plus ``prod(...)`` + ``apply(..., batch_size=[...])``.
td_ndim = tensordict.ndim
if td_ndim == 0:
raise ValueError(
"LSTMModule(recurrent_mode=True) requires the input "
"tensordict to have at least one batch dim (time). Got a "
"0-d tensordict."
)
elif td_ndim == 1:
tensordict_shaped = tensordict.unsqueeze(0)
elif td_ndim == 2:
tensordict_shaped = tensordict
else:
tensordict_shaped = tensordict.flatten(0, -2)
else:
tensordict_shaped = tensordict.reshape(-1).unsqueeze(-1)

is_init = tensordict_shaped["is_init"].squeeze(-1)
splits = None
backend = self.recurrent_backend
if backend == "auto":
# In eager, CuDNN-backed pad is the fastest path; under torch.compile
# the data-dependent ``_split_and_pad_sequence`` branch is unfriendly,
# so prefer scan there.
backend = "scan" if is_compiling() else "pad"
use_scan = self.recurrent_mode and backend == "scan"
use_triton = self.recurrent_mode and backend == "triton"
Expand All @@ -948,12 +954,10 @@ def forward(self, tensordict: TensorDictBase):
):
from torchrl.objectives.value.utils import _get_num_per_traj_init

# if we have consecutive trajectories, things get a little more complicated
# we have a tensordict of shape [B, T]
# we will split / pad things such that we get a tensordict of shape
# [N, T'] where T' <= T and N >= B is the new batch size, such that
# each index of N is an independent trajectory. We'll need to keep
# track of the indices though, as we want to put things back together in the end.
# Multi-trajectory rollouts under the pad backend: split each row
# into per-trajectory windows of shape [N, T'], run the LSTM on
# the padded result, then stitch them back. Required for correctness
# whenever ``is_init`` fires mid-row.
splits = _get_num_per_traj_init(is_init)
tensordict_shaped_shape = tensordict_shaped.shape
tensordict_shaped = _split_and_pad_sequence(
Expand Down Expand Up @@ -999,7 +1003,6 @@ def forward(self, tensordict: TensorDictBase):
tensordict_shaped.set(self.out_keys[1], hidden0)
tensordict_shaped.set(self.out_keys[2], hidden1)
if splits is not None:
# let's recover our original shape
tensordict_shaped = _inv_pad_sequence(tensordict_shaped, splits).reshape(
tensordict_shaped_shape
)
Expand Down Expand Up @@ -2120,21 +2123,21 @@ def forward(self, tensordict: TensorDictBase):
# we want to get an error if the value input is missing, but not the hidden states
defaults = [NO_DEFAULT, None]
shape = tensordict.shape
tensordict_shaped = tensordict
if self.recurrent_mode:
# if less than 2 dims, unsqueeze
ndim = tensordict_shaped.get(self.in_keys[0]).ndim
while ndim < 3:
tensordict_shaped = tensordict_shaped.unsqueeze(0)
ndim += 1
if ndim > 3:
dims_to_flatten = ndim - 3
# we assume that the tensordict can be flattened like this
nelts = prod(tensordict_shaped.shape[: dims_to_flatten + 1])
tensordict_shaped = tensordict_shaped.apply(
lambda value: value.flatten(0, dims_to_flatten),
batch_size=[nelts, tensordict_shaped.shape[-1]],
# Straight-line shape normalization (see LSTMModule.forward).
td_ndim = tensordict.ndim
if td_ndim == 0:
raise ValueError(
"GRUModule(recurrent_mode=True) requires the input "
"tensordict to have at least one batch dim (time). Got a "
"0-d tensordict."
)
elif td_ndim == 1:
tensordict_shaped = tensordict.unsqueeze(0)
elif td_ndim == 2:
tensordict_shaped = tensordict
else:
tensordict_shaped = tensordict.flatten(0, -2)
else:
tensordict_shaped = tensordict.reshape(-1).unsqueeze(-1)

Expand All @@ -2153,12 +2156,6 @@ def forward(self, tensordict: TensorDictBase):
):
from torchrl.objectives.value.utils import _get_num_per_traj_init

# if we have consecutive trajectories, things get a little more complicated
# we have a tensordict of shape [B, T]
# we will split / pad things such that we get a tensordict of shape
# [N, T'] where T' <= T and N >= B is the new batch size, such that
# each index of N is an independent trajectory. We'll need to keep
# track of the indices though, as we want to put things back together in the end.
splits = _get_num_per_traj_init(is_init)
tensordict_shaped_shape = tensordict_shaped.shape
tensordict_shaped = _split_and_pad_sequence(
Expand Down Expand Up @@ -2190,7 +2187,6 @@ def forward(self, tensordict: TensorDictBase):
tensordict_shaped.set(self.out_keys[0], val)
tensordict_shaped.set(self.out_keys[1], hidden)
if splits is not None:
# let's recover our original shape
tensordict_shaped = _inv_pad_sequence(tensordict_shaped, splits).reshape(
tensordict_shaped_shape
)
Expand Down
Loading