For Torch, I have come across Named Tensors, which should work in a similar way: https://docs.pytorch.org/docs/stable/named_tensor.html

The docs say that it's a prototype feature, and I think it has been that way for a few years now, so no idea how production-ready it is.

It's a much worse API than Xarrays, it seems like somebody should build it on top of PyTorch.

I really like einops. This works for numpy, pytorch and keras/tensorflow and has easy named transpose, repeat, and eimsum operations.

Same - I’ve been using einops and jaxtyping together pretty extensively recently and it helps a lot for reading/writing multidimensional array code. Also array_api_compat, the API coverage isn’t perfect but it’s pretty satisfying to write code that works for both PyTorch and numpy arrays

https://docs.kidger.site/jaxtyping/

https://data-apis.org/array-api-compat/

For pytorch the analogue is Named Tensors, but it's a provisional feature and not supported everywhere.

https://docs.pytorch.org/docs/stable/named_tensor.html