memory profile etc

[cnsgsz]

Is anyone able to see anything memory_profile (under tools dropdown menu) for on tensorboard? I always see "There is no memory profile to display because there were no memory activity data in the captured duration."

I saw doc on memory profiling with pprof. Is it recommended to use that instead of tensorboard?

On a related note, it seems that both tensorboard and perfetto supports traces (doc). Is either one recommended over the other? Does either one able to show tf data traces along side jax traces?

[skye]

For memory profiling, check out the memory_viewer tab in Tensorboard. The Cloud TPU docs have a good description of this tab, and I believe it should work on all platforms, not just TPU. We should update the JAX profiling page to mention this!

For Tensorboard vs. Perfetto, they both show exactly the same trace events. Perfetto is often quicker to get going with since the trace is hosted (no installation, remote port forwarding, etc.), but Tensorboard lets you open an old trace and has other information beyond the execution trace (e.g. the memory_viewer tab). So up to you, they're both good options depending on what you prefer.