pytorch named tensors <> einops integration

[arogozhnikov]

pytorch named tensors <> einops integration

[arogozhnikov]

Hello Daniel,

thanks for bringing this up. I had similar thoughts when reading pytorch documentation of named tensors.

Simplest thing to try is:

• check non-decomposed features on the input, fail when axis name does not coincide with one in einops
• set result's axes names when result axis is not a composed axis

Examples:

rearrange(x, 'b (c1 c2) h w -> b c1 h w') # second axis name will not be checked, will fail if for any other axis name was set, but did not coincide
rearrange(x, 'b (c1 c2) h w -> h w c1 c2 b') # result names are exactly h, w, c1, c2, and b
rearrange(x, 'b (c1 c2) h w -> h (c1 c2) w b ()') # result names are exactly h, None, w, b and None

We can try to add this as an experimental module like einops.named_pytorch with separate versions of rearrange, reduce and layers.

[arogozhnikov]

I think this repo may be of interest as an experiment of interaction einops <-> xarray (xarray names axes),
https://github.com/arviz-devs/xarray-einstats

I don't have plans to integrate with pytorch's named_tensor; projects related to einops <-> named tensors are welcome in discussions.