並行處理

JAX 對 Python 並行處理的支援有限。

用戶端可以從不同的 Python 執行緒並行呼叫 JAX API（例如，jit() 或 grad()）。

不允許從多個執行緒並行操作 JAX 追蹤值。換句話說，雖然可以從多個執行緒呼叫使用 JAX 追蹤的函式（例如，jit()），但您不得使用執行緒來操作傳遞給 jit() 的函式 f 實作內部的 JAX 值。如果您這樣做，最有可能的結果是 JAX 傳回不明錯誤。