jax.profiler.start_trace

jax.profiler.start_trace(log_dir, create_perfetto_link=False, create_perfetto_trace=False)

Starts a profiler trace.

The trace will capture CPU, GPU, and/or TPU activity, including Python functions and JAX on-device operations. Use stop_trace() to end the trace and save the results to log_dir.

The resulting trace can be viewed with TensorBoard. Note that TensorBoard doesn’t need to be running when collecting the trace.

Only one trace may be collected at a time. A RuntimeError will be raised if start_trace() is called while another trace is running.

Parameters:

• log_dir – The directory to save the profiler trace to (usually the TensorBoard log directory).

• create_perfetto_link (bool) – A boolean which, if true, creates and prints link to the Perfetto trace viewer UI (https://ui.perfetto.dev). The program will block until the link is opened and Perfetto loads the trace.

• create_perfetto_trace (bool) – A boolean which, if true, additionally dumps a perfetto_trace.json.gz file that is compatible for upload with the Perfetto trace viewer UI (https://ui.perfetto.dev). The file will also be generated if create_perfetto_link is true. This could be useful if you want to generate a Perfetto-compatible trace without blocking the process.

Return type:

None