[BugFix] Pointwise ops with a batch-shaped tensor give the right values for a nested lazy stack

tensordict/base.py

@@ -666,7 +666,11 @@ def _same_leaf_keys(td: TensorDictBase, others: tuple) -> bool:
 
 
 def _apply_pointwise(
-    td: TensorDictBase, fn: Callable, others: tuple, inplace: bool
+    td: TensorDictBase,
+    fn: Callable,
+    others: tuple,
+    inplace: bool,
+    is_leaf: Callable[[Type], bool] | None = None,
 ) -> TensorDictBase:
     """Computes ``fn(leaf, *others)`` leaf by leaf, pairing the leaves of the tensordicts in ``others`` by key."""
     others = tuple(map(_unwrap_tensorclass, others))
@@ -680,14 +684,14 @@ def leaf_fn(x, *leaves):
 
     tds = [other for flag, other in zip(is_td, others) if flag]
     if inplace:
-        return td.apply_(leaf_fn, *tds)
+        return td.apply_(leaf_fn, *tds, is_leaf=is_leaf)
     # apply ignores the keys that only the others have; the dense ops raise
     keys = set(td.keys(True, True))
     for other in tds:
         other_keys = set(other.keys(True, True))
         if other_keys != keys:
             _mismatch_keys(keys, other_keys)
-    return td.apply(leaf_fn, *tds, propagate_lock=True)
+    return td.apply(leaf_fn, *tds, propagate_lock=True, is_leaf=is_leaf)
 
 
 def _maybe_broadcast_other(op: str, n_other: int = 1) -> Callable[[Callable], Callable]:
@@ -764,6 +768,18 @@ def new_func(self, *others, **kwargs):
             if any(isinstance(other, torch.Tensor) for other in others_map):
                 from tensordict._lazy import LazyStackedTensorDict
 
+                wrapped_attr = type(self_expand)._wrapped_attr
+                if wrapped_attr is not None:
+                    source = getattr(self_expand, wrapped_attr)
+                    if isinstance(source, LazyStackedTensorDict) and not is_non_tensor(
+                        source
+                    ):
+                        # A TypedTensorDict on a lazy stack: the op on the stack
+                        # gives each member its slice of the operands.
+                        result = getattr(source, op)(*others_map, *args, **kwargs)
+                        if inplace and self_expand is self:
+                            return self
+                        return type(self_expand)._wrap_td(result)
                 if isinstance(self_expand, LazyStackedTensorDict) and not is_non_tensor(
                     self_expand
                 ):
@@ -789,18 +805,17 @@ def new_func(self, *others, **kwargs):
                         stack_dim=stack_dim,
                         stack_dim_name=self_expand._td_dim_name,
                     )
+                # A lazy stack below the root would run the op on the leaves of
+                # its members, which lack the stack dim: call the op on the
+                # stack, or on the tensorclass or TypedTensorDict that wraps it.
+                is_leaf = (
+                    _is_leaf_or_lazy_stack if _has_lazy_stack(self_expand) else None
+                )
                 if (
                     n_other > 1
                     and any(_is_tensor_collection(type(o)) for o in others_map)
                     # an empty tensordict returns an empty result, as before
                     and not self_expand.is_empty()
-                    # the leaves of a nested lazy stack lack the stack dim, so
-                    # expand_as_right would misalign the tensors: raise as before
-                    and not any(
-                        _has_nested_lazy_stack(o)
-                        for o in (self_expand, *others_map)
-                        if _is_tensor_collection(type(o))
-                    )
                 ):
                     if op == "clamp":
                         # clamp rejects a tensordict bound next to a tensor one
@@ -822,6 +837,7 @@ def new_func(self, *others, **kwargs):
                         ),
                         others_map,
                         inplace,
+                        is_leaf=is_leaf,
                     )
                 result = self_expand._fast_apply(
                     lambda x: leaf_op(
@@ -832,7 +848,8 @@ def new_func(self, *others, **kwargs):
                         ],
                         *args,
                         **kwargs,
-                    )
+                    ),
+                    is_leaf=is_leaf,
                 )
                 # Like torch, an in-place op returns self once it has written into
                 # the tensors that self holds. Containers that build their entries
@@ -1245,6 +1262,26 @@ def _default_is_leaf(cls: Type) -> bool:
     return not _is_tensor_collection(cls) or getattr(cls, "_pass_through", False)
 
 
+def _is_leaf_or_lazy_stack(cls: Type) -> bool:
+    """Returns ``True`` for a leaf of ``_default_is_leaf``, a lazy stack of tensordicts, and a tensorclass or TypedTensorDict, which may wrap one.
+
+    The pointwise ops call themselves on a nested lazy stack, which gives each
+    member its slice of a batch-shaped tensor: the leaves of a member lack the
+    stack dim (#3005).
+    """
+    if _default_is_leaf(cls):
+        return True
+    if _is_non_tensor(cls):
+        return False
+    from tensordict._lazy import LazyStackedTensorDict
+
+    return (
+        issubclass(cls, LazyStackedTensorDict)
+        or _is_tensorclass(cls)
+        or getattr(cls, "_wrapped_attr", None) is not None
+    )
+
+
 def _flatten_keys_leaves(
     td: TensorCollection, is_leaf: Callable[[Type], bool]
 ) -> list[tuple[NestedKey, Any]]:
@@ -1894,22 +1931,6 @@ def _has_lazy_stack(td: TensorDictBase) -> bool:
     )
 
 
-def _has_nested_lazy_stack(td: TensorDictBase) -> bool:
-    """Whether a tensordict below the root of ``td``, or below the members of a lazy stack ``td``, is a lazy stack."""
-    from tensordict._lazy import LazyStackedTensorDict
-
-    td = _unwrap(td)
-    if isinstance(td, LazyStackedTensorDict):
-        return any(_has_lazy_stack(member) for member in td.tensordicts)
-    if td._reads_return_copies:
-        return False
-    return any(
-        _has_lazy_stack(value)
-        for value in td.values()
-        if _is_tensor_collection(type(value)) and not _is_non_tensor(type(value))
-    )
-
-
 def _all_leaf_keys(td: TensorDictBase) -> list:
     """The keys of the leaves of ``td``, non-tensor ones included, as tuples.
 
@@ -8753,6 +8774,7 @@ def func(*args, **kwargs):
                 is_leaf is None
                 or is_leaf is _default_is_leaf
                 or is_leaf is _NESTED_TENSORS_AS_LISTS
+                or is_leaf is _is_leaf_or_lazy_stack
             )
             and self._is_dense_node()
         ):

test/tensordict/test_lazy.py

@@ -1092,6 +1092,83 @@ def test_arithmetic_ops_broadcast_tensor(self, stack_dim):
             assert lazy.add_(other) is lazy
             assert (lazy == dense + other).all()
 
+    @pytest.mark.parametrize("stack_dim", [0, 1])
+    @pytest.mark.parametrize("depth", [1, 2])
+    def test_arithmetic_ops_broadcast_tensor_nested(self, stack_dim, depth):
+        # A lazy stack nested in a tensordict also gives each member its slice
+        # of a batch-shaped tensor. Here the leaves of a member have the shape
+        # of the batch, so a misaligned tensor would give wrong values (#3005)
+        def make(lazy):
+            n = TensorDict(a=torch.arange(8.0).view(2, 2, 2), batch_size=[2, 2])
+            if lazy:
+                n = lazy_stack(list(n.unbind(stack_dim)), stack_dim)
+            td = TensorDict(n=n, b=torch.arange(4.0).view(2, 2), batch_size=[2, 2])
+            if depth == 2:
+                td = TensorDict(o=td, batch_size=[2, 2])
+            return td
+
+        def check(result, expected):
+            assert set(result.keys(True, True)) == set(expected.keys(True, True))
+            for key, value in expected.items(True, True):
+                torch.testing.assert_close(result[key], value)
+
+        lazy_key = ("n",) if depth == 1 else ("o", "n")
+        for other in (torch.arange(4.0).view(2, 2) * 10, torch.tensor([1.0, 2.0])):
+            for op in ("add", "sub", "mul", "maximum", "clamp_min"):
+                result = getattr(make(True), op)(other)
+                assert isinstance(result.get(lazy_key), LazyStackedTensorDict)
+                check(result, getattr(make(False), op)(other))
+            check(make(True) + other, make(False) + other)
+            check(
+                make(True).lerp(make(True) + 1, other),
+                make(False).lerp(make(False) + 1, other),
+            )
+            check(
+                make(True).addcmul(other, make(True) + 1),
+                make(False).addcmul(other, make(False) + 1),
+            )
+            td = make(True)
+            assert td.mul_(other) is td
+            check(td, make(False) * other)
+
+    @pytest.mark.parametrize("stack_dim", [0, 1])
+    @pytest.mark.parametrize("kind", ["tensorclass", "typed"])
+    def test_arithmetic_ops_broadcast_tensor_nested_wrapped(self, stack_dim, kind):
+        # A tensorclass or a TypedTensorDict that wraps a lazy stack, nested in
+        # a tensordict, also gives each member its slice of a batch-shaped
+        # tensor (#3005)
+        class MC(TensorClass):
+            a: torch.Tensor
+
+        class Typed(TypedTensorDict):
+            a: torch.Tensor
+
+        def make(lazy):
+            n = TensorDict(a=torch.arange(8.0).view(2, 2, 2), batch_size=[2, 2])
+            if lazy:
+                n = lazy_stack(list(n.unbind(stack_dim)), stack_dim)
+            n = (MC if kind == "tensorclass" else Typed).from_tensordict(n)
+            return TensorDict(n=n, batch_size=[2, 2])
+
+        def check(result, expected):
+            assert type(result["n"]) is type(expected["n"])
+            torch.testing.assert_close(result["n", "a"], expected["n", "a"])
+
+        other = torch.arange(4.0).view(2, 2) * 10
+        for op in ("add", "mul", "maximum"):
+            check(getattr(make(True), op)(other), getattr(make(False), op)(other))
+        check(
+            make(True).lerp(make(True) + 1, other),
+            make(False).lerp(make(False) + 1, other),
+        )
+        check(
+            make(True).addcmul(other, make(True) + 1),
+            make(False).addcmul(other, make(False) + 1),
+        )
+        td = make(True)
+        assert td.add_(other) is td
+        check(td, make(False) + other)
+
     @pytest.mark.parametrize("stack_dim", [0, 1])
     @pytest.mark.parametrize(
         "op,args",

test/tensordict/test_misc.py

@@ -789,22 +789,25 @@ class TC:
     )
     @pytest.mark.parametrize("tensor_first", [False, True])
     def test_nested_lazy_and_tensor_operands(self, op, tensor_first):
-        # With a lazy stack nested in the tensordict, the leaves of each member
-        # lack the stack dim and a batch-shaped tensor would be misaligned, so
-        # the ops still raise
-        def make_td():
-            return TensorDict(
-                n=LazyStackedTensorDict(
-                    TensorDict(a=torch.arange(2.0)), TensorDict(a=torch.arange(2.0))
-                ),
-                batch_size=[2],
-            )
-
-        args = [make_td() + 10, torch.tensor([0.0, 1.0])]
-        if tensor_first:
-            args = args[::-1]
-        with pytest.raises(TypeError):
-            getattr(make_td(), op)(*args)
+        # With a lazy stack nested in the tensordict, each member gets its
+        # slice of a batch-shaped tensor, as for a dense tensordict (#3005)
+        def make_td(lazy):
+            n = TensorDict(a=torch.arange(6.0).view(2, 3) + 1, batch_size=[2, 3])
+            if lazy:
+                n = LazyStackedTensorDict(*n.unbind(0))
+            return TensorDict(n=n, batch_size=[2, 3])
+
+        def make_args(lazy):
+            args = [make_td(lazy) + 10, torch.arange(6.0).view(2, 3) / 6]
+            return args[::-1] if tensor_first else args
+
+        td = make_td(True)
+        result = getattr(td, op)(*make_args(True))
+        expected = getattr(make_td(False), op)(*make_args(False))
+        assert isinstance(result["n"], LazyStackedTensorDict)
+        torch.testing.assert_close(result["n", "a"], expected["n", "a"])
+        if op.endswith("_"):
+            assert result is td
 
     @pytest.mark.parametrize("op", ["maximum", "minimum"])
     @pytest.mark.parametrize("other_type", ["tensordict", "float", "tensor"])