jax.default_device

jax.default_device = <jax._src.config.State object>

用於 jax_default_device 配置選項的上下文管理器。

設定 JAX 運算的預設裝置。設定為裝置物件 (例如 jax.devices("cpu")[0]) 以使用該裝置作為 JAX 運算和 jit 編譯函式呼叫的預設裝置 (對多裝置計算沒有影響，例如 pmapped 函式呼叫)。設定為 None 以使用系統預設裝置。請參閱 控制裝置上的資料和計算放置 以取得有關裝置放置的更多資訊。

參數:

new_val (Any)