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 格式的协议缓冲区 (protocol buffer) 捕获 JAX 设备内存分析数据。
save_device_memory_profile(filename[, backend])	收集设备内存分析数据并将其写入文件。