Whether/how to add functionality (and create pull requests)

[sunsibar]

Is there a way to make more torch functions work for TensorDicts? For example, I'm missing torch.repeat_interleave(my_tensor_dict, ...) . How would I go about implementing this myself, as in, what file to look in for examples how other torch-compatible functionality is implemented?

[vmoens]

Thanks for the question,
There are two use cases:

• You want to apply torch.repeat_interleave along a feature dimension (that is not part of the batch-size). In that case, you should be calling tensordict.apply(lambda tensor: torch.repeat_interleave(tensor, ...))
• You want to apply torch.repeat_interleave along a batch dimension. Then you can code that yourself (can require a bit of experience) or make a request. For us, coming up with a new function like this is fairly easy and can be done in a couple of hours...

[sunsibar]

Thanks for the quick reply! My use case is the second, I was asking where I would start coding that, so where I would generally implement new features to then make a pull request. Does there already exist a source file implementing other torch-like functions?

I've now just added a helper function that does what I need; the question here was about how to contribute, for someone not experienced with your code base.

[vmoens]

We don't have recipies (yet!) but that's a great idea.
Here's how it generally works:

• Usually we will want to have a function that works recursively, either on tensors or on tensordict items. A plain tensordict.apply(func) will only be called on the tensors, which does not fit the purpose. (If you want to understand why this approach is better, think of a TensorDictBase subclass -- call it MyTensorDict -- with a dedicated torch.repeat_interleave. If we call a plain apply, we won't be treating this class accurately if it's nested within a regular TensorDict).
• Instead, we want to call TensorDict._fast_apply(func, call_on_nested=True). This will call your func on tensors and tensordicts, regardeless of their type. (taking the example above, if we place a MyTensorDict within a TensorDict, we can be sure that the repeat_interleave of MyTensorDict will be associated with it as it should).
• It is likely that func will change the batch size. Fortunately, apply (and therefore _fast_apply) allow you to override the batch-size.
• Once we know all of that, we have to think about what func will look like. Well, func will simply be torch.repeat_interleave itself! How can we make this happen? PyTorch offers a way to augment torch functions to new types. To make this work with TensorDict, we just need to do something like

@implement_for_td(torch.repeat_interleave)
def _repeat_interleave(tensordict, repeats, dim=None, *, output_size=None):
    new_batch_size = ... # compute the new batch size
    if dim < 0:
        # if we have negative dim, we will mess things up when calling `repeat_interleave` on tensors or nested tensordicts
        # with more dimensions than the root
        dim = tensordict.ndim + dim
    return tensordict._fast_apply(
        lambda x: torch.repeat_interleave(x, repeats=repeats, dim=dim, *, output_size=output_size),
        batch_size=new_batch_size,
    )

What happens in that function is that we will iterate over the values of tensordict one at a time. Regardless of the value type (tensor or tensordict), torch.repeat_interleave will be called. If it's a tensordict, this function (_repeat_interleave) will be called, otherwise the native torch.repeat_interleave will be used.

This should be placed in _torch_func.py

You should also test that it works across all TensorDictBase subtypes using the TestTensorDicts class:

class TestTensorDicts(TestTensorDictsBase):
    ...
    @pytest.mark.parametrize("dim,repeats", [[-1, repeats0], [-2, repeats1], [0, repeats2], [1, repeats3]]
    def test_repeat_interleave(self, td_name, device, dim, repeats):
         tensordict = getattr(self, td_name)(device)
         rtd = torch.repeat_interleave(tensordict, repeats, dim)
         ... # do some checks

I'll leave it to you to figure out what dim and repeats values should be!

If you want examples, you can have a look at _view, _reshape, _transpose or _permute in _td.py