Is there anything similar to this for something like Tensorflow, Keras or Pytorch? I haven't used them super recently, but in the past I needed to do all of the things you just described in painful to debug ways.
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
For pytorch the analogue is Named Tensors, but it's a provisional feature and not supported everywhere.