性能分析# 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 – 如果为 true，则启用命名调用包装，用于标记性能分析跟踪。（请参阅 enabled_named_call）。
