性能分析

flax.linen.enable_named_call()

启用命名调用包装以标记性能分析跟踪。

启用命名调用包装后，模块中执行的所有 JAX 操作都将在 jax.named_scope 下运行。 Module 类名将显示在 Tensorboard 性能分析 UI 中属于该模块的操作周围，从而简化性能分析过程。

请注意，jax.named_scope 仅适用于已编译的函数（例如：使用 jax.jit 或 jax.pmap）。

flax.linen.disable_named_call()

禁用命名调用包装。

参见 enable_named_call

flax.linen.override_named_call(enable=True)

返回一个上下文管理器，用于启用/禁用命名调用包装。

参数

enable – 如果为真，则启用命名调用包装以标记性能分析跟踪。（参见 enabled_named_call）。