筆記

本節包含關於使用 JAX 相關主題的簡短筆記；另請參閱 JAX 增強提案 (JEPs) 中更長的設計討論。

相依性與版本相容性

• API 相容性 概述了 JAX 關於跨版本 API 相容性的政策。

• Python 和 NumPy 版本支援政策 概述了 JAX 關於與 Python 和 NumPy 相容性的政策。

遷移與棄用

• jax.Array 遷移 總結了 jax v 0.4.1 中預設陣列型別的變更

記憶體與運算使用量

• 非同步分派 描述了 JAX 的非同步分派模型。

• 並行性 描述了 JAX 如何與其他 Python 並行性互動。

• GPU 記憶體分配 描述了 JAX 如何與 GPU 上的記憶體分配互動。

程式設計師注意事項

• 階數提升警告 描述如何配置 jax.numpy 以避免隱含的階數提升。