剖析

flax.linen.enable_named_call()

針對標記剖析追蹤啟用命名呼叫包裝。

當命名呼叫包裝已啟用時，在 `Module` 中執行的所有 `JAX` 操作都將在 `jax.named_scope` 中運行。`Module` 類別名稱將顯示在屬於 `Tensorboard` 剖析 UI 中該 `Module` 的操作周圍，簡化剖析流程。

請注意，`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`）。