jax.profiler 模块

追踪和时间性能分析

性能分析计算 描述了如何使用 JAX 的追踪和时间性能分析功能。

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

设备内存性能分析

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

device_memory_profile([后端])	将 JAX 设备内存性能分析捕获为 pprof 格式的协议缓冲区。
save_device_memory_profile(文件名[, 后端])	收集设备内存性能分析并将其写入文件。