笔记

本节包含有关使用 JAX 的相关主题的简要笔记；另请参阅 JAX 增强提案 (JEPs) 中的更长设计讨论。

依赖项和版本兼容性

• API 兼容性 概述了 JAX 在发布版本之间的 API 兼容性策略。

• Python 和 NumPy 版本支持政策 概述了 JAX 与 Python 和 NumPy 的兼容性策略。

内存和计算使用

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

• 并发 描述了 JAX 如何与其他 Python 并发交互。

• GPU 内存分配 描述了 JAX 如何与 GPU 上的内存分配交互。

程序员防护栏

• 秩提升警告 描述了如何配置 jax.numpy 以避免隐式秩提升。

数组和数据类型

• 类型提升语义 描述了 JAX 对两个或多个值的函数的隐式类型提升。

• 默认数据类型和 X64 标志 描述了 JAX 如何确定数组创建函数的默认数据类型。