jax.profiler 模块

跟踪和时间剖析

剖析计算 描述了如何利用 JAX 的跟踪和时间剖析功能。

start_server(port)	在 port 端口上启动剖析器服务器。
start_trace(log_dir[, create_perfetto_link, ...])	启动一个剖析器跟踪。
stop_trace()	停止当前正在运行的剖析器跟踪。
trace(log_dir[, create_perfetto_link, ...])	用于进行剖析器跟踪的上下文管理器。
annotate_function(func[, name])	为函数执行生成跟踪事件的装饰器。
TraceAnnotation(*args, **kwargs)	在剖析器中生成跟踪事件的上下文管理器。
StepTraceAnnotation(name, **kwargs)	在剖析器中生成步进跟踪事件的上下文管理器。

设备内存剖析

有关 JAX 设备内存剖析功能的介绍，请参阅 剖析设备内存。

device_memory_profile([backend])	以 pprof 格式的协议缓冲区捕获 JAX 设备内存剖析。
save_device_memory_profile(filename[, backend])	收集设备内存剖析并将其写入文件。