JAX 中文文档(十七)

简介: JAX 中文文档(十七)


原文:jax.readthedocs.io/en/latest/

JAX 术语表

原文:jax.readthedocs.io/en/latest/glossary.html

数组

JAX 的 numpy.ndarray 的类比。见 jax.Array

CPU

缩写Central Processing Unit,CPU 是大多数计算机中可用的标准计算架构。JAX 可以在 CPU 上运行计算,但通常在 GPU 和 TPU 上可以实现更好的性能。

设备

用于指代 JAX 用于执行计算的 CPU、GPU 或 TPU 的通用名称。

forward-mode autodiff

见 JVP

函数式编程

一种编程范式,程序通过应用和组合纯函数定义。JAX 设计用于函数式程序。

GPU

缩写Graphical Processing Unit,GPU 最初专门用于图像渲染相关操作,但现在更通用。JAX 能够针对 GPU 进行快速数组操作(参见 CPU 和 TPU)。

jaxpr

缩写JAX Expression,jaxpr 是由 JAX 生成的计算的中间表示形式,转发到 XLA 进行编译和执行。详见 Understanding Jaxprs 以获取更多讨论和示例。

JIT

缩写Just In Time 编译,JIT 在 JAX 中通常指将数组操作编译为 XLA,通常使用 jax.jit() 完成。

JVP

缩写Jacobian Vector Product,有时也称为forward-mode 自动微分。有关详细信息,请参阅 Jacobian-Vector products (JVPs, aka forward-mode autodiff)。在 JAX 中,JVP 是通过 jax.jvp() 实现的转换。另见 VJP。

primitive

primitive 是 JAX 程序中使用的基本计算单位。jax.lax 中的大多数函数代表单个原语。在 jaxpr 中表示计算时,jaxpr 中的每个操作都是一个原语。

纯函数

纯函数是仅基于其输入生成输出且没有副作用的函数。JAX 的转换模型设计用于处理纯函数。参见 functional programming。

pytree

pytree 是一个抽象,允许 JAX 以统一的方式处理元组、列表、字典和其他更一般的包含数组值的容器。请参阅 Working with pytrees 以获取更详细的讨论。

reverse-mode autodiff

见 VJP。

SPMD

缩写Single Program Multi Data,指的是一种并行计算技术,即在不同设备(例如几个 TPU)上并行运行相同计算(例如神经网络的前向传播)对不同输入数据(例如批处理中的不同输入)的计算。jax.pmap() 是实现 SPMD 并行性的 JAX 转换。

static

在 JIT 编译中,未被追踪的值(参见 Tracer)。有时也指静态值的编译时计算。

TPU

张量处理单元 的缩写,TPUs 是专门为深度学习应用中的 N 维张量快速运算而设计的芯片。JAX 能够针对 TPUs 进行快速数组操作(另见 CPU 和 GPU)。

追踪器

一个用作 JAX 数组替身的对象,以确定 Python 函数执行的操作序列。在内部,JAX 通过 jax.core.Tracer 类实现此功能。

转换

高阶函数:即接受函数作为输入并输出转换后函数的函数。在 JAX 中的示例包括 jax.jit()jax.vmap()jax.grad()

VJP

向量雅可比积,有时也称为反向模式自动微分。有关详细信息,请参阅向量雅可比积(VJPs,又称反向模式自动微分)。在 JAX 中,VJP 是通过 jax.vjp() 实现的转换。还请参阅 JVP。

XLA

加速线性代数 的缩写,XLA 是一个专用于线性代数操作的编译器,是 JIT 编译 JAX 代码的主要后端。请参阅 www.tensorflow.org/xla/

弱类型

JAX 数据类型,其类型提升语义与 Python 标量相同;请参阅 JAX 中的弱类型值。

相关实践学习
部署Stable Diffusion玩转AI绘画(GPU云服务器)
本实验通过在ECS上从零开始部署Stable Diffusion来进行AI绘画创作,开启AIGC盲盒。
相关文章
|
4月前
|
存储 API 索引
JAX 中文文档(十五)(5)
JAX 中文文档(十五)
53 3
|
4月前
|
机器学习/深度学习 存储 API
JAX 中文文档(十五)(4)
JAX 中文文档(十五)
40 3
|
4月前
|
机器学习/深度学习 数据可视化 编译器
JAX 中文文档(十四)(5)
JAX 中文文档(十四)
51 2
|
4月前
JAX 中文文档(十一)(5)
JAX 中文文档(十一)
17 1
|
4月前
|
API 异构计算 Python
JAX 中文文档(十一)(4)
JAX 中文文档(十一)
33 1
|
4月前
|
算法 API 开发工具
JAX 中文文档(十二)(5)
JAX 中文文档(十二)
55 1
|
4月前
|
机器学习/深度学习 存储 API
JAX 中文文档(十五)(2)
JAX 中文文档(十五)
36 0
|
4月前
|
安全 API 网络架构
JAX 中文文档(十五)(1)
JAX 中文文档(十五)
51 0
|
4月前
|
TensorFlow API 算法框架/工具
JAX 中文文档(十五)(3)
JAX 中文文档(十五)
29 0
|
4月前
|
API 异构计算 索引
JAX 中文文档(十四)(2)
JAX 中文文档(十四)
44 0