- 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
The resulting trace can be viewed with TensorBoard. Note that TensorBoard doesn’t need to be running when collecting the trace.
Only once trace may be collected a time. A RuntimeError will be raised if
start_trace()is called while another trace is running.
log_dir – The directory to save the profiler trace to (usually the TensorBoard log directory).
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.
bool) – A boolean which, if true, additionally dumps a
perfetto_trace.json.gzfile that is compatible for upload with the Perfetto trace viewer UI (https://ui.perfetto.dev). The file will also be generated if
create_perfetto_linkis true. This could be useful if you want to generate a Perfetto-compatible trace without blocking the processs.