[Feature] Allow neg dims during LazyStack

ghstack-source-id: dd2a9bd
Pull Request resolved: #1240

Author: Vincent Moens

tensordict/_lazy.py

@@ -250,9 +250,13 @@ def __init__(
             td0 = tensordicts[0]
             device = td0.device
         if stack_dim < 0:
-            raise RuntimeError(
-                f"stack_dim must be non negative, got stack_dim={stack_dim}"
-            )
+            ndim = td0.ndim
+            try:
+                stack_dim = _maybe_correct_neg_dim(stack_dim, ndim=ndim + 1, shape=None)
+            except Exception:
+                raise RuntimeError(
+                    f"Couldn't infer stack dim from negative value, got stack_dim={stack_dim}"
+                )
         _batch_size = td0.batch_size
         if stack_dim > len(_batch_size):
             raise RuntimeError(

test/test_tensordict.py

@@ -9395,6 +9395,19 @@ def test_lazy_stacked_insert(self, dim, index, device):
         with pytest.raises(ValueError, match="Batch sizes in tensordicts differs"):
             lstd.insert(index, TensorDict({"a": torch.ones(17)}, [17], device=device))
 
+    def test_neg_dim_lazystack(self):
+        td0 = TensorDict(batch_size=(3, 5))
+        td1 = TensorDict(batch_size=(4, 5))
+        assert lazy_stack([td0, td1], -1).shape == (-1, 5, 2)
+        assert lazy_stack([td0, td1], -2).shape == (-1, 2, 5)
+        assert lazy_stack([td0, td1], -3).shape == (2, -1, 5)
+        with pytest.raises(RuntimeError):
+            assert lazy_stack([td0, td1], -4)
+        with pytest.raises(RuntimeError):
+            assert lazy_stack([td0, td1, TensorDict()])
+        with pytest.raises(RuntimeError):
+            assert lazy_stack([TensorDict(), td0, td1])
+
     @pytest.mark.parametrize(
         "reduction", ["sum", "nansum", "mean", "nanmean", "std", "var", "prod"]
     )