A really cool post and a great set of visualizations!
Computing sparse Jacobians can save a lot of compute if there's a real lack of dependency between part of the input and the output. Discovering this automatically through coloring is very appealing.
Another alternative is to implement sparse rules for each operation yourself, but that often requires custom autodiff implementations which aren't easy to get right, I wrote a small toy version of a sparse rules-based autodiff here: https://github.com/rdyro/SpAutoDiff.jl
nvtop can actually support TPUs too via https://github.com/rdyro/libtpuinfo/ https://github.com/Syllo/nvtop/blob/76890233d759199f50ad3bdb...