术语表

术语表#

数组 (Array)#: JAX 中与 numpy.ndarray 类似的结构。请参阅 jax.Array。
CPU#: 中央处理器的缩写，CPU 是大多数计算机中可用的标准计算架构。JAX 可以在 CPU 上运行计算，但通常在 GPU 和 TPU 上可以获得更好的性能。
设备 (Device)#: 用于指代 JAX 用于执行计算的 CPU、GPU 或 TPU 的通用名称。
前向模式自动微分#: 请参阅 JVP
函数式编程#: 一种编程范式，其中程序通过应用和组合纯函数来定义。JAX 被设计为与函数式程序一起使用。
GPU#: 图形处理单元 (Graphical Processing Unit) 的缩写。GPU 最初专门用于屏幕图像渲染相关操作，但现在其用途更加广泛。JAX 能够利用 GPU 对数组进行快速操作（另请参阅CPU 和 TPU）。
jaxpr#: JAX 表达式 (JAX expression) 的缩写。jaxpr 是 JAX 生成的计算的中间表示形式，并被转发到 XLA 进行编译和执行。有关更多讨论和示例，请参阅JAX 内部：jaxpr 语言。
JIT#: 即时 (Just In Time) 编译的缩写。在 JAX 中，JIT 通常指将数组操作编译为 XLA，通常使用 jax.jit() 完成。
JVP#: 雅可比向量积 (Jacobian Vector Product) 的缩写，有时也称为前向模式自动微分。有关更多详细信息，请参阅雅可比-向量积（JVP，又名正向模式自动微分）。在 JAX 中，JVP 是一个通过 jax.jvp() 实现的转换。另请参阅 VJP。
原语#: 原语是 JAX 程序中使用的基本计算单元。jax.lax 中的大多数函数都表示单个原语。在 jaxpr 中表示计算时，jaxpr 中的每个操作都是一个原语。
纯函数#: 纯函数是指其输出仅基于其输入，并且没有副作用的函数。JAX 的转换模型旨在与纯函数一起使用。另请参阅函数式编程。
pytree#: pytree 是一种抽象，它允许 JAX 以统一的方式处理元组、列表、字典以及其他更通用的数组值容器。有关更详细的讨论，请参阅使用 pytree。
反向模式自动微分#: 请参阅 VJP。
SPMD#: 单程序多数据 (Single Program Multi Data) 的缩写，它指的是一种并行计算技术，其中相同的计算（例如，神经网络的前向传递）在不同设备（例如，多个 TPU）上并行运行在不同的输入数据（例如，批量中的不同输入）上。jax.pmap() 是一个实现 SPMD 并行的 JAX 转换。
静态#: 在 JIT 编译中，一个未被追踪的值（参见 Tracer）。有时也指对静态值进行编译时计算。
TPU#: 张量处理单元 (Tensor Processing Unit) 的缩写。TPU 是专门为深度学习应用中 N 维张量的快速操作而设计的芯片。JAX 能够利用 TPU 对数组进行快速操作（另请参阅CPU 和 GPU）。
Tracer#: 一个用作 JAX Array 的占位符的对象，用于确定 Python 函数执行的操作序列。在内部，JAX 通过 jax.core.Tracer 类实现这一点。
转换#: 一个高阶函数：即，一个以函数作为输入并输出转换后的函数的函数。JAX 中的示例包括 jax.jit()、jax.vmap() 和 jax.grad()。
VJP#: 向量雅可比积 (Vector Jacobian Product) 的缩写，有时也称为反向模式自动微分。有关更多详细信息，请参阅向量-雅可比积（VJP，又名反向模式自动微分）。在 JAX 中，VJP 是一个通过 jax.vjp() 实现的转换。另请参阅 JVP。
XLA#: 加速线性代数 (Accelerated Linear Algebra) 的缩写。XLA 是一种用于线性代数操作的特定领域编译器，它是 JIT 编译的 JAX 代码的主要后端。请参阅 https://tensorflowcn.cn/xla/。
弱类型#: 一种 JAX 数据类型，它具有与 Python 标量相同的类型提升语义；请参阅JAX 中的弱类型值。