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).