JAX: High performance array computing — JAX documentation

JAX 官方文档网站整体介绍
该网站是JAX（一款面向加速器的数组计算与程序转换Python库）的官方文档平台，聚焦高性能数值计算与大规模机器学习场景，提供从入门到进阶的完整指导，帮助开发者理解和使用JAX的核心功能与生态工具。

核心内容模块与功能
1. 快速入门引导
基础准备：提供详细的安装指南，覆盖不同环境下的配置步骤；
入门教程：包含“Quickstart: How to think in JAX”（JAX思维快速入门）、“JAX The Sharp Bits”（JAX常见陷阱），帮用户快速掌握JAX的核心逻辑与避坑技巧；
JAX 101系列：系统讲解JAX的关键能力，包括即时编译（JIT）、自动向量化、自动微分、pytrees处理、伪随机数生成、并行编程入门、JIT下的控制流与逻辑运算符、追踪（Tracing）、有状态计算等内容。

2. 资源、指南与参考
核心概念：解析JAX的底层设计逻辑（如数组模型、函数转换等）；
高级指南：覆盖分布式计算（分布式数组、显式分片、手动并行、多控制器JAX）、性能优化（GPU性能提示、编译缓存、内存分析）、调试工具（编译打印/断点、checkify转换、调试标志）、自动微分进阶（自定义导数规则、梯度检查点）、导出与序列化（跨框架交互如JAX2TF、形状多态）、内核语言Pallas（TPU/GPU内核编写、流水线处理、矩阵乘法优化）等；
API参考：详细列出JAX的公共API模块，包括`jax.numpy`（兼容NumPy的数组操作）、`jax.scipy`（科学计算工具）、`jax.lax`（底层线性代数操作）、`jax.random`（随机数生成）、`jax.sharding`（分片管理）、`jax.debug`（调试工具）等，每个模块下的函数均有具体说明（如`jax.numpy.fft.fft`、`jax.scipy.stats.bernoulli.logpmf`）；
实验性模块：包含`jax.experimental`下的前沿功能，如稀疏数组操作（`bcoo`格式）、Pallas内核语言（TPU/GPU高性能计算）、编译缓存、自定义分区等。

3. 开发者相关
贡献指南：说明如何向JAX项目贡献代码，包括从源码构建、 regression问题排查等；
JAX增强提案（JEPs）：展示JAX的设计演进（如PRNG设计、自定义导数规则、版本控制、类型提升语义等）；
Autodidax系列：从底层实现角度讲解JAX核心逻辑（如Autodidax: JAX core from scratch）。

4. 扩展与生态
扩展指南：教用户如何基于JAX构建自定义功能（如自定义Jaxpr解释器、使用`jax.extend`模块）；
生态系统：列出基于JAX的第三方工具，覆盖神经网络（Flax、Equinox、Keras）、优化器（Optax、Optimistix）、数据加载（Grain、TensorFlow Datasets）、概率编程（Blackjax、Numpyro）、物理模拟（JAX MD、Brax）、大语言模型（MaxText、AXLearn）等方向，并提供社区维护的“Awesome JAX”列表。

特色亮点
API友好：采用NumPy风格API，降低开发者学习成本；
多后端支持：代码可在CPU、GPU、TPU上无缝运行；
可组合转换：函数转换（编译、自动微分等）可叠加使用，灵活性高；
生态丰富：覆盖机器学习、科学计算、模拟等多个领域，提供完整的工具链支持。