备注

本节包含关于使用 JAX 相关主题的简短备注；另请参阅 JAX 增强提案 (JEPs) 中的更长篇幅的设计讨论。

依赖项和版本兼容性

• API 兼容性 概述了 JAX 在跨版本 API 兼容性方面的策略。

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

内存和计算使用

• 异步调度 描述了 JAX 的异步调度模型。

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

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

程序员防护措施

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

数组和数据类型

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

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