start_server(port)

Starts the profiler server on portport.

start_trace(log_dir[, create_perfetto_link, ...])

Starts a profiler trace.

stop_trace()

Stops the currently-running profiler trace.

trace(log_dir[, create_perfetto_link, ...])

Context manager to take a profiler trace.

annotate_function(func[, name])

Decorator that generates a trace event for the execution of a function.

TraceAnnotation(*args, **kwargs)

Context manager that generates a trace event in the profiler.

StepTraceAnnotation(name, **kwargs)

Context manager that generates a step trace event in the profiler.

device_memory_profile([backend])

Captures a JAX device memory profile aspprof-format protocol buffer.

save_device_memory_profile(filename[, backend])

Collects a device memory profile and writes it to a file.