[BUG] TypedTensorDict reshape, masked_select and repeat_interleave with inplace=True return a new object instead of self #2494

inplace=True promises to keep the object identity, and code that relies on it, such as assert td.reshape(6, inplace=True) is td or holding a reference across the call, breaks for TypedTensorDict. The TypedTensorDict is reshaped in place, but reshape, masked_select and repeat_interleave return a second wrapper around the same data. flatten, repeat, gather and pad return self as expected.

import torch
from tensordict import TypedTensorDict

class State(TypedTensorDict):
    x: torch.Tensor

for name, op in [
    ("flatten", lambda s: s.flatten(0, 1, inplace=True)),
    ("reshape", lambda s: s.reshape(6, inplace=True)),
    ("masked_select", lambda s: s.masked_select(torch.ones(2, 3, dtype=torch.bool), inplace=True)),
    ("repeat_interleave", lambda s: s.repeat_interleave(2, dim=0, inplace=True)),
]:
    s = State(x=torch.zeros(2, 3), batch_size=[2, 3])
    print(name, op(s) is s)

flatten True
reshape False
masked_select False
repeat_interleave False

Expected: True for every line, as on a dense TensorDict, where td.reshape(6, inplace=True) is td holds.

Cause: reshape, masked_select and repeat_interleave are in _WRAP_DELEGATES (tensordict/typedtensordict.py:1015-1030). Their delegate (_make_delegate_wrap, tensordict/typedtensordict.py:277-288) calls the method on self._source, which changes the source in place and returns it, and then wraps that result in a new instance with type(self)._wrap_td(result). The methods that return self are not overridden, so they run the shared implementation on the TypedTensorDict itself.

Possible fix: in _make_delegate_wrap, return self when result is self._source. Only these in-place calls change; the out-of-place calls still return a new wrapper.

tensordict main (eba7879), torch 2.14.1, Python 3.12.12.