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.
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.