I've never really understood the point of these visualizer things. The idea that a model is always well represented by a directed acyclic graph seems extremely dated.<p>I really would love a PyTorch/JAX profiler that shows, in annotated Python, where your code is allocating memory, using compute or doing device copies.