基于 JAX 构建#
学习 JAX 高级用法的一个绝佳方法是观察其他库如何使用 JAX：它们如何将该库集成到自己的 API 中，它在数学上增加了哪些功能，以及它是如何被用于实现其他库的计算加速的。
以下示例展示了如何利用 JAX 的功能在众多领域和软件包中定义加速计算。
梯度计算#
简便的梯度计算是 JAX 的核心功能。在 JaxOpt 库 中，value 和 grad 被直接用于其源代码中的多种优化算法中，以供用户调用。
同样，上述提到的 Dynamax 与 Optax 的搭配使用，便是通过梯度实现历史上难以完成的估算方法的例子，例如 使用 Optax 进行最大似然估计。
在跨多设备的单核上实现计算加速#
在 JAX 中定义的模型可以通过 JIT 编译来实现单次计算加速。编译后的代码通常无需任何更改，即可发送到 CPU、GPU 或 TPU 设备以获得额外加速。这使得从开发到生产的工作流程能够顺利衔接。在 Dynamax 中，线性状态空间模型求解器中计算开销最大的部分已被 jit 化。另一个更复杂的例子来自 PyTensor，它动态编译 JAX 函数并随后 jit 化所构建的函数。
使用并行化实现单机和多机加速#
JAX 的另一个优势是使用 pmap 和 vmap 函数调用或装饰器来并行化计算非常简单。在 Dynamax 中，状态空间模型通过 VMAP 装饰器 进行并行化，多目标跟踪便是一个实际应用案例。
将 JAX 代码整合到您或您用户的工作流程中#
JAX 具有高度的可组合性，可以通过多种方式使用。JAX 可以采用独立模式使用，即用户自行定义所有计算。然而，也可以采用其他模式，例如使用基于 JAX 构建并提供特定功能的库。这些库可以定义特定类型的模型（如神经网络或状态空间模型），或提供特定的功能（如优化）。以下是每种模式的更具体示例。
直接使用#
JAX 可以直接导入并用于“从零开始”构建模型，如本网站各处所示，例如 JAX 101 教程 或 使用 JAX 的神经网络。如果您无法找到针对特定问题的现成代码，或者希望减少代码库中的依赖项数量，这可能是最佳选择。
使用 JAX 构建可组合的特定领域库#
另一种常见方法是使用提供预构建功能的包，无论是模型定义还是某种类型的计算。这些包的组合可以混合使用，以实现完整的端到端工作流程，即定义模型并估算其参数。
一个例子是 Flax，它简化了神经网络的构建。Flax 通常与 Optax 搭配使用，其中 Flax 定义神经网络架构，而 Optax 提供优化和模型拟合功能。
另一个是 Dynamax，它允许轻松定义状态空间模型。在 Dynamax 中，参数可以使用 使用 Optax 的最大似然估计 进行估算，或者可以使用 来自 Blackjax 的 MCMC 估算完整的贝叶斯后验分布。