[Feature suggestion] Support of tensordict (and tensorclass)

[fzimmermann89]

Tensordicts are dicts of tensors with a common batch dimension, tensorclasses dataclasses of tensor.

https://pytorch.org/tensordict/

Tensordict seems to be still not fully mature, but is already quite useful. If I am not mistaken, it supports most operations required to write a einops backend -- einsum is missing.

Before I (or somebody else) does the work to draft an implementation, I am wondering if this might be included in einops or if @arogozhnikov would either like to wait and see if tensordict gets merged into pytorch proper (pytorch/pytorch#112441) or if the missing einsum functionality would be a stopper.

For us, it would already be immensely helpful to be able to use rearrange on the batch dimensions of tensordicts.

Cheers

[arogozhnikov]

Hi Felix,

missing einsum functionality would be a stopper.

not really. From brief look, min/max are also missing.

• I guess it would 'just work' with einops with a simple backend
• It is a no-go to depend on in-development lib in einops; torchdict isn't compatible with the way other libs are tested
• array api won't work for this case

My suggestion: implement a backend (see TorchBackend for example), put it in a separate file. When you import that file, einops will pick up your backend.

If there are other cases of non-standard, but 'likely working' backends, I can think of mechanism to auto-import those.

[arogozhnikov]

also, don't forget to update this issue if it worked or if there are other hurdles

[vmoens]

We added min/max/amin/amax.

The upstreaming to core is still roadmapped.

Is there anything I can do on the tensordict side to unblock this?