Google Xprof
Xprof is an open-source performance profiling and visualization tool developed under the OpenXLA umbrella at Google. It provides deep insights into the execution of ML workloads on accelerators (TPUs, GPUs, CPUs) for XLA-based frameworks including JAX, TensorFlow, and PyTorch/XLA.
I contributed to this project as an open-source contributor.