From e69f12cf36886a2757cf00b03d659e9928edef15 Mon Sep 17 00:00:00 2001 From: Vincent Moens Date: Fri, 15 May 2026 21:47:44 +0100 Subject: [PATCH] Update [ghstack-poisoned] --- torchrl/modules/tensordict_module/rnn.py | 153 ++++++++--------------- 1 file changed, 52 insertions(+), 101 deletions(-) diff --git a/torchrl/modules/tensordict_module/rnn.py b/torchrl/modules/tensordict_module/rnn.py index e63c98dd9bf..712ceba8b14 100644 --- a/torchrl/modules/tensordict_module/rnn.py +++ b/torchrl/modules/tensordict_module/rnn.py @@ -15,16 +15,11 @@ 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 -from torchrl._utils import ( - _ContextManager, - _DecoratorContextManager, - implement_for, - is_compiling, -) +from torchrl._utils import _ContextManager, _DecoratorContextManager, implement_for from torchrl.data.tensor_specs import Unbounded # ``torch._higher_order_ops.scan`` was introduced in PyTorch 2.6. Gate the @@ -654,7 +649,7 @@ def __init__( proj_size=0, bidirectional=False, python_based=False, - recurrent_backend: typing.Literal["auto", "pad", "scan", "triton"] = "pad", + recurrent_backend: typing.Literal["auto", "pad", "scan", "triton"] = "scan", recurrent_compute_dtype: torch.dtype = torch.float32, *, in_key=None, @@ -748,6 +743,12 @@ def __init__( self.in_keys = in_keys self.out_keys = out_keys self._recurrent_mode = default_recurrent_mode + # Resolve "auto" once at construction. Previously this was resolved + # per-forward via is_compiling(), which produced different backend + # choices in eager vs compiled traces and could trigger recompiles + # when entering/leaving a compiled region. + if recurrent_backend == "auto": + recurrent_backend = "scan" self.recurrent_backend = recurrent_backend self.recurrent_compute_dtype = recurrent_compute_dtype @@ -907,59 +908,40 @@ def set_recurrent_mode(self, mode: bool = True): @dispatch def forward(self, tensordict: TensorDictBase): - from torchrl.objectives.value.functional import ( - _inv_pad_sequence, - _split_and_pad_sequence, - ) - # 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. + # No Python loop on tensor ndim, no data-dependent branches. + 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 is resolved at construction time; no runtime is_compiling() + # check, no data-dependent fallback. The "pad" backend in + # recurrent_mode no longer handles trajectory resets — callers with + # multi-trajectory rollouts must use 'scan' or 'triton'. Use + # 'pad' only for single-trajectory windows or non-recurrent-mode + # collection. backend = self.recurrent_backend - if backend == "auto": - backend = "scan" if is_compiling() else "pad" use_scan = self.recurrent_mode and backend == "scan" use_triton = self.recurrent_mode and backend == "triton" - if ( - self.recurrent_mode - and not use_scan - and not use_triton - and is_init[..., 1:].any() - ): - 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( - tensordict_shaped.select(*self.in_keys, strict=False), splits - ) - is_init = tensordict_shaped["is_init"].squeeze(-1) + splits = None value, hidden0, hidden1 = ( tensordict_shaped.get(key, default) @@ -998,11 +980,6 @@ def forward(self, tensordict: TensorDictBase): tensordict_shaped.set(self.out_keys[0], val) 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 - ) if shape != tensordict_shaped.shape or tensordict_shaped is not tensordict: tensordict.update(tensordict_shaped.reshape(shape)) @@ -1865,7 +1842,7 @@ def __init__( dropout=0, bidirectional=False, python_based=False, - recurrent_backend: typing.Literal["auto", "pad", "scan", "triton"] = "pad", + recurrent_backend: typing.Literal["auto", "pad", "scan", "triton"] = "scan", recurrent_compute_dtype: torch.dtype = torch.float32, *, in_key=None, @@ -1956,6 +1933,10 @@ def __init__( self.in_keys = in_keys self.out_keys = out_keys self._recurrent_mode = default_recurrent_mode + # Resolve "auto" once at construction. See LSTMModule.__init__ for + # rationale. + if recurrent_backend == "auto": + recurrent_backend = "scan" self.recurrent_backend = recurrent_backend self.recurrent_compute_dtype = recurrent_compute_dtype @@ -2112,59 +2093,34 @@ def set_recurrent_mode(self, mode: bool = True): @dispatch @set_lazy_legacy(False) def forward(self, tensordict: TensorDictBase): - from torchrl.objectives.value.functional import ( - _inv_pad_sequence, - _split_and_pad_sequence, - ) - # 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) is_init = tensordict_shaped["is_init"].squeeze(-1) - splits = None + # Backend resolved at construction. See LSTMModule.forward for the + # rationale on dropping the pad-with-resets fallback. backend = self.recurrent_backend - if backend == "auto": - backend = "scan" if is_compiling() else "pad" use_scan = self.recurrent_mode and backend == "scan" use_triton = self.recurrent_mode and backend == "triton" - if ( - self.recurrent_mode - and not use_scan - and not use_triton - and is_init[..., 1:].any() - ): - 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( - tensordict_shaped.select(*self.in_keys, strict=False), splits - ) - is_init = tensordict_shaped["is_init"].squeeze(-1) + splits = None value, hidden = ( tensordict_shaped.get(key, default) @@ -2189,11 +2145,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 - ) if shape != tensordict_shaped.shape or tensordict_shaped is not tensordict: tensordict.update(tensordict_shaped.reshape(shape))