[Feature Request] Support propagation of in-place modifications to masked tensordict to the parent tensordict.

[benoitdescamps]

Motivation

For tensors, the following apply,

x = torch.randn(3,4)
mask = [True, False, True]
y = torch.randn(2,4)
x[mask] = y

torch.testing.assert_allclose(x[mask],y)

on the other hand for tensordict

    td = TensorDict({"a": torch.randn(3, 4, 2), "b": torch.randn(3, 4)}, [3, 4])

    mask = torch.tensor([True, False, True])
    x = torch.randn(2, 4, 2)
    td[mask]["a"] = x
    torch.testing.assert_allclose(td[mask]["a"], x)

returns error AssertionError: Tensor-likes are not close!

Solution

Instead of return a new Tensordict using mask-select, we can return instead a SubTensorDict.

Checklist

• I have checked that there is no similar issue in the repo (required)