jax.lax 模块#

jax.lax 是一个底层原语操作库,它支撑了诸如 jax.numpy 之类的库。转换规则,例如 JVP 和批处理规则,通常被定义为对 jax.lax 原语的转换。

许多原语是围绕等效 XLA 操作的简单包装器,由 XLA 操作语义 文档描述。在少数情况下,JAX 与 XLA 存在差异,通常是为了确保操作集在 JVP 和转置规则的操作下是封闭的。

尽可能使用诸如 jax.numpy 的库,而不是直接使用 jax.laxjax.numpy 的 API 遵循 NumPy,因此比 jax.lax 的 API 更稳定,且不太可能发生变化。

运算符#

abs(x)

逐元素绝对值:\(|x|\)

acos(x)

逐元素反余弦:\(\mathrm{acos}(x)\)

acosh(x)

逐元素反双曲余弦:\(\mathrm{acosh}(x)\)

add(x, y)

逐元素加法:\(x + y\)

after_all(*operands)

合并一个或多个 XLA 令牌值。

approx_max_k(operand, k[, ...])

以近似方式返回 operand 中最大的 k 个值及其索引。

approx_min_k(operand, k[, ...])

以近似方式返回 operand 中最小的 k 个值及其索引。

argmax(operand, axis, index_dtype)

计算沿 axis 的最大元素的索引。

argmin(operand, axis, index_dtype)

计算沿 axis 的最小元素的索引。

asin(x)

逐元素反正弦:\(\mathrm{asin}(x)\)

asinh(x)

逐元素反双曲正弦:\(\mathrm{asinh}(x)\)

atan(x)

逐元素反正切:\(\mathrm{atan}(x)\)

atan2(x, y)

逐元素两个变量的反正切:\(\mathrm{atan}({x \over y})\)

atanh(x)

逐元素反双曲正切:\(\mathrm{atanh}(x)\)

batch_matmul(lhs, rhs[, precision])

批量矩阵乘法。

bessel_i0e(x)

指数缩放的0阶修正贝塞尔函数:\(\mathrm{i0e}(x) = e^{-|x|} \mathrm{i0}(x)\)

bessel_i1e(x)

指数缩放的1阶修正贝塞尔函数:\(\mathrm{i1e}(x) = e^{-|x|} \mathrm{i1}(x)\)

betainc(a, b, x)

逐元素正则化不完全贝塔积分。

bitcast_convert_type(operand, new_dtype)

逐元素位强制转换。

bitwise_and(x, y)

逐元素与:\(x \wedge y\)

bitwise_not(x)

逐元素非:\(\neg x\)

bitwise_or(x, y)

逐元素或:\(x \vee y\)

bitwise_xor(x, y)

逐元素异或:\(x \oplus y\)

population_count(x)

逐元素位计数,计算每个元素中设置位的数量。

broadcast(operand, sizes)

广播数组,添加新的前导维度

broadcast_in_dim(operand, shape, ...)

包装 XLA 的 BroadcastInDim 运算符。

broadcast_shapes()

返回 NumPy 广播 shapes 后得到的结果形状。

broadcast_to_rank(x, rank)

添加前导维度 1 以使 x 的秩为 rank

broadcasted_iota(dtype, shape, dimension)

围绕 iota 的便捷包装器。

cbrt(x)

逐元素立方根:\(\sqrt[3]{x}\)

ceil(x)

逐元素向上取整:\(\left\lceil x \right\rceil\)

clamp(min, x, max)

逐元素钳位。

clz(x)

逐元素前导零计数。

collapse(operand, start_dimension[, ...])

将数组的维度折叠成单个维度。

complex(x, y)

逐元素创建复数:\(x + jy\)

concatenate(operands, dimension)

沿着 dimension 连接一系列数组。

conj(x)

逐元素复共轭函数:\(\overline{x}\)

conv(lhs, rhs, window_strides, padding[, ...])

围绕 conv_general_dilated 的便捷包装器。

convert_element_type(operand, new_dtype)

逐元素强制转换。

conv_dimension_numbers(lhs_shape, rhs_shape, ...)

将卷积 dimension_numbers 转换为 ConvDimensionNumbers

conv_general_dilated(lhs, rhs, ...[, ...])

一般的 n 维卷积运算符,可选扩张。

conv_general_dilated_local(lhs, rhs, ...[, ...])

具有可选扩张的一般 n 维非共享卷积运算符。

conv_general_dilated_patches(lhs, ...[, ...])

提取受 conv_general_dilated 感受野影响的补丁。

conv_transpose(lhs, rhs, strides, padding[, ...])

计算 N 维卷积“转置”的便捷包装器。

conv_with_general_padding(lhs, rhs, ...[, ...])

围绕 conv_general_dilated 的便捷包装器。

cos(x)

逐元素余弦:\(\mathrm{cos}(x)\)

cosh(x)

逐元素双曲余弦:\(\mathrm{cosh}(x)\)

cumlogsumexp(operand[, axis, reverse])

沿 axis 计算累积 logsumexp。

cummax(operand[, axis, reverse])

沿 axis 计算累积最大值。

cummin(operand[, axis, reverse])

沿 axis 计算累积最小值。

cumprod(operand[, axis, reverse])

沿 axis 计算累积积。

cumsum(operand[, axis, reverse])

沿 axis 计算累积和。

digamma(x)

逐元素 digamma 函数:\(\psi(x)\)

div(x, y)

逐元素除法:\(x \over y\)

dot(lhs, rhs[, precision, ...])

向量/向量、矩阵/向量和矩阵/矩阵乘法。

dot_general(lhs, rhs, dimension_numbers[, ...])

通用点积/收缩运算符。

dynamic_index_in_dim(operand, index[, axis, ...])

围绕 dynamic_slice 的便捷包装器,用于执行整数索引。

dynamic_slice(operand, start_indices, ...)

包装 XLA 的 DynamicSlice 运算符。

dynamic_slice_in_dim(operand, start_index, ...)

围绕 lax.dynamic_slice() 的便捷包装器,应用于一个维度。

dynamic_update_index_in_dim(operand, update, ...)

围绕 dynamic_update_slice() 的便捷包装器,用于更新单个 axis 中大小为 1 的切片。

dynamic_update_slice(operand, update, ...)

包装 XLA 的 DynamicUpdateSlice 运算符。

dynamic_update_slice_in_dim(operand, update, ...)

围绕 dynamic_update_slice() 的便捷包装器,用于更新单个 axis 中的切片。

eq(x, y)

逐元素等于:\(x = y\)

erf(x)

逐元素误差函数:\(\mathrm{erf}(x)\)

erfc(x)

逐元素互补误差函数:\(\mathrm{erfc}(x) = 1 - \mathrm{erf}(x)\)

erf_inv(x)

逐元素误差函数的反函数:\(\mathrm{erf}^{-1}(x)\)

exp(x)

逐元素指数函数:\(e^x\)

expand_dims(array, dimensions)

在数组中插入任意数量的大小为 1 的维度。

expm1(x)

逐元素计算 \(e^{x} - 1\)

fft(x, fft_type, fft_lengths)

floor(x)

逐元素向下取整:\(\left\lfloor x \right\rfloor\)

full(shape, fill_value[, dtype, sharding])

返回一个形状为 shape、填充值为 fill_value 的数组。

full_like(x, fill_value[, dtype, shape, ...])

基于示例数组 x 创建一个类似于 np.full 的填充数组。

gather(operand, start_indices, ...[, ...])

收集运算符。

ge(x, y)

逐元素大于等于:\(x \geq y\)

gt(x, y)

逐元素大于:\(x > y\)

igamma(a, x)

逐元素正则化不完全伽马函数。

igammac(a, x)

逐元素互补正则化不完全伽马函数。

imag(x)

逐元素提取虚部:\(\mathrm{Im}(x)\)

index_in_dim(operand, index[, axis, keepdims])

围绕 lax.slice() 的便捷包装器,用于执行整数索引。

index_take(src, idxs, axes)

integer_pow(x, y)

逐元素幂运算:\(x^y\),其中 \(y\) 是一个固定的整数。

iota(dtype, size)

包装 XLA 的 Iota 运算符。

is_finite(x)

逐元素 \(\mathrm{isfinite}\)

le(x, y)

逐元素小于等于:\(x \leq y\)

lgamma(x)

逐元素对数伽马函数:\(\mathrm{log}(\Gamma(x))\)

log(x)

逐元素自然对数:\(\mathrm{log}(x)\)

log1p(x)

逐元素计算 \(\mathrm{log}(1 + x)\)

logistic(x)

逐元素逻辑斯蒂函数(sigmoid 函数):\(\frac{1}{1 + e^{-x}}\)

lt(x, y)

逐元素小于:\(x < y\)

max(x, y)

逐元素最大值:\(\mathrm{max}(x, y)\)

min(x, y)

逐元素最小值:\(\mathrm{min}(x, y)\)

mul(x, y)

逐元素乘法:\(x \times y\)

ne(x, y)

逐元素不等于:\(x \neq y\)

neg(x)

逐元素取负:\(-x\)

nextafter(x1, x2)

返回在x2方向上x1之后的下一个可表示值。

optimization_barrier(operand, /)

阻止编译器跨越屏障移动操作。

pad(operand, padding_value, padding_config)

对数组应用低、高和/或内部填充。

platform_dependent(*args[, default])

分阶段执行特定于平台的代码。

polygamma(m, x)

逐元素多伽马函数:\(\psi^{(m)}(x)\)

population_count(x)

逐元素位计数,计算每个元素中设置位的数量。

pow(x, y)

逐元素幂运算:\(x^y\)

random_gamma_grad(a, x)

来自Gamma(a, 1)的样本的逐元素导数。

real(x)

逐元素提取实部:\(\mathrm{Re}(x)\)

reciprocal(x)

逐元素倒数:\(1 \over x\)

reduce(operands, init_values, computation, ...)

封装 XLA 的 Reduce 运算符。

reduce_precision(operand, exponent_bits, ...)

封装 XLA 的 ReducePrecision 运算符。

reduce_window(operand, init_value, ...[, ...])

rem(x, y)

逐元素取余:\(x \bmod y\)

reshape(operand, new_sizes[, dimensions])

封装 XLA 的 Reshape 运算符。

rev(operand, dimensions)

封装 XLA 的 Rev 运算符。

rng_bit_generator(key, shape[, dtype, algorithm])

无状态 PRNG 位生成器。

rng_uniform(a, b, shape)

有状态 PRNG 生成器。

round(x[, rounding_method])

逐元素舍入。

rsqrt(x)

逐元素倒数平方根:\(1 \over \sqrt{x}\)

scatter(operand, scatter_indices, updates, ...)

散射更新运算符。

scatter_add(operand, scatter_indices, ...[, ...])

散射加法运算符。

scatter_apply(operand, scatter_indices, ...)

散射应用运算符。

scatter_max(operand, scatter_indices, ...[, ...])

散射最大值运算符。

scatter_min(operand, scatter_indices, ...[, ...])

散射最小值运算符。

scatter_mul(operand, scatter_indices, ...[, ...])

散射乘法运算符。

shift_left(x, y)

逐元素左移:\(x \ll y\)

shift_right_arithmetic(x, y)

逐元素算术右移:\(x \gg y\)

shift_right_logical(x, y)

逐元素逻辑右移:\(x \gg y\)

sign(x)

逐元素符号函数。

sin(x)

逐元素正弦函数:\(\mathrm{sin}(x)\)

sinh(x)

逐元素双曲正弦函数:\(\mathrm{sinh}(x)\)

slice(operand, start_indices, limit_indices)

封装 XLA 的 Slice 运算符。

slice_in_dim(operand, start_index, limit_index)

围绕 lax.slice() 的便捷包装器,仅应用于一个维度。

排序()

封装 XLA 的 Sort 运算符。

sort_key_val(keys, values[, dimension, ...])

沿着 dimensionkeys 进行排序,并将相同的排列应用于 values

sqrt(x)

逐元素平方根:\(\sqrt{x}\)

square(x)

逐元素平方:\(x^2\)

squeeze(array, dimensions)

从数组中压缩任意数量的大小为 1 的维度。

sub(x, y)

逐元素减法:\(x - y\)

tan(x)

逐元素正切函数:\(\mathrm{tan}(x)\)

tanh(x)

逐元素双曲正切函数:\(\mathrm{tanh}(x)\)

top_k(operand, k)

返回 operand 最后一个轴上的前 k 个值及其索引。

transpose(operand, permutation)

封装 XLA 的 Transpose 运算符。

zeros_like_array(x)

zeta(x, q)

逐元素 Hurwitz zeta 函数:\(\zeta(x, q)\)

控制流运算符#

associative_scan(fn, elems[, reverse, axis])

并行执行具有结合二元运算的扫描。

cond(pred, true_fun, false_fun, *operands[, ...])

有条件地应用 true_funfalse_fun

fori_loop(lower, upper, body_fun, init_val, *)

lowerupper 的循环,通过降维到 jax.lax.while_loop() 实现。

map(f, xs, *[, batch_size])

在数组的前导轴上映射函数。

scan(f, init[, xs, length, reverse, unroll, ...])

在数组的前导轴上扫描函数,同时传递状态。

select(pred, on_true, on_false)

根据布尔谓词在两个分支之间进行选择。

select_n(which, *cases)

从多个案例中选择数组值。

switch(index, branches, *operands[, operand])

根据 index 应用恰好一个 branches

while_loop(cond_fun, body_fun, init_val)

cond_fun 为 True 时,重复调用 body_fun 进行循环。

自定义梯度算子#

stop_gradient(x)

停止梯度计算。

custom_linear_solve(matvec, b, solve[, ...])

使用隐式定义的梯度执行无矩阵线性求解。

custom_root(f, initial_guess, solve, ...[, ...])

可微地求解函数的根。

并行算子#

all_gather(x, axis_name, *[, ...])

收集所有副本中 x 的值。

all_to_all(x, axis_name, split_axis, ...[, ...])

具体化映射轴并映射不同的轴。

psum(x, axis_name, *[, axis_index_groups])

在 pmapped 轴 axis_name 上计算 x 的全归约和。

psum_scatter(x, axis_name, *[, ...])

类似于 psum(x, axis_name),但每个设备仅保留部分结果。

pmax(x, axis_name, *[, axis_index_groups])

在 pmapped 轴 axis_name 上计算 x 的全归约最大值。

pmin(x, axis_name, *[, axis_index_groups])

在 pmapped 轴 axis_name 上计算 x 的全归约最小值。

pmean(x, axis_name, *[, axis_index_groups])

在 pmapped 轴 axis_name 上计算 x 的全归约平均值。

ppermute(x, axis_name, perm)

根据排列 perm 执行集体排列。

pshuffle(x, axis_name, perm)

jax.lax.ppermute 的便捷包装器,使用备用排列编码。

pswapaxes(x, axis_name, axis, *[, ...])

将 pmapped 轴 axis_name 与未映射轴 axis 交换。

axis_index(axis_name)

返回映射轴 axis_name 上的索引。

线性代数算子 (jax.lax.linalg)#

cholesky(x, *[, symmetrize_input])

Cholesky 分解。

eig(x, *[, compute_left_eigenvectors, ...])

一般矩阵的特征分解。

eigh(x, *[, lower, symmetrize_input, ...])

厄米特矩阵的特征分解。

hessenberg(a)

将方阵简化为上 Hessenberg 形式。

lu(x)

带部分主元的 LU 分解。

householder_product(a, taus)

基本 Householder 反射器的乘积。

qdwh(x, *[, is_hermitian, max_iterations, ...])

基于 QR 的动态加权 Halley 迭代,用于极分解。

qr(x, *[, full_matrices])

QR 分解。

schur(x, *[, compute_schur_vectors, ...])

svd()

奇异值分解。

triangular_solve(a, b, *[, left_side, ...])

三角求解。

tridiagonal(a, *[, lower])

将对称/厄米特矩阵简化为三对角形式。

tridiagonal_solve(dl, d, du, b)

计算三对角线性系统的解。

参数类#

class jax.lax.ConvDimensionNumbers(lhs_spec, rhs_spec, out_spec)[source]#

描述卷积的批次、空间和特征维度。

参数:
  • lhs_spec (Sequence[int]) – 包含 (批次维度,特征维度,空间维度…) 的非负整数维度编号元组。

  • rhs_spec (Sequence[int]) – 包含 (输出特征维度,输入特征维度,空间维度…) 的非负整数维度编号元组。

  • out_spec (Sequence[int]) – 包含 (批次维度,特征维度,空间维度…) 的非负整数维度编号元组。

jax.lax.ConvGeneralDilatedDimensionNumbers#

tuple[str, str, str] 的别名 | ConvDimensionNumbers | None

class jax.lax.GatherDimensionNumbers(offset_dims, collapsed_slice_dims, start_index_map)[source]#

描述 XLA 的 Gather 运算符 的维度编号参数。有关维度编号含义的更多详细信息,请参阅 XLA 文档。

参数:
  • offset_dims (tuple[int, ...]) – gather 输出中偏移到从 operand 切片得到的数组的一组维度。必须是按升序排列的整数元组,每个整数表示输出的维度编号。

  • collapsed_slice_dims (tuple[int, ...]) – operand 中满足 slice_sizes[i] == 1 且不应在 gather 输出中具有对应维度的一组维度 i。必须是按升序排列的整数元组。

  • start_index_map (tuple[int, ...]) – 对于 start_indices 中的每个维度,给出要从中切片的 operand 中的对应维度。必须是大小等于 start_indices.shape[-1] 的整数元组。

与 XLA 的 GatherDimensionNumbers 结构不同,index_vector_dim 是隐式的;始终存在索引向量维度,并且它必须始终是最后一个维度。要收集标量索引,请添加一个大小为 1 的尾随维度。

class jax.lax.GatherScatterMode(value)[source]#

描述如何在 gather 或 scatter 中处理越界索引。

可能的值是

CLIP

索引将被钳位到最近的范围内值,即,使得要收集的整个窗口都在范围内。

FILL_OR_DROP

如果收集的窗口的任何部分超出范围,则返回的整个窗口(即使其他元素在范围内)都将填充一个常数。如果散布的窗口的任何部分超出范围,则整个窗口将被丢弃。

PROMISE_IN_BOUNDS

用户保证索引在范围内。不会执行其他检查。实际上,使用当前的 XLA 实现,这意味着越界 gather 将被钳位,但越界 scatter 将被丢弃。如果索引超出范围,则梯度将不正确。

class jax.lax.Precision(value)[source]#

lax 矩阵乘法相关函数的精度枚举。

JAX 函数中与设备相关的 precision 参数通常控制加速器后端(即 TPU 和 GPU)上数组计算的速度和精度之间的权衡。对 CPU 后端没有影响。这仅对 float32 计算有效,并且不影响输入/输出数据类型。成员是

DEFAULT

最快的模式,但精度最低。在 TPU 上:以 bfloat16 执行 float32 计算。在 GPU 上:如果可用,则使用 tensorfloat32(例如在 A100 和 H100 GPU 上),否则使用标准 float32(例如在 V100 GPU 上)。别名:'default''fastest'

HIGH

速度较慢但精度较高。在 TPU 上:在 3 个 bfloat16 传递中执行 float32 计算。在 GPU 上:在可用情况下使用 tensorfloat32,否则使用 float32。别名:'high'..

HIGHEST

速度最慢但精度最高。在 TPU 上:在 6 个 bfloat16 中执行 float32 计算。别名:'highest'。在 GPU 上:使用 float32。

jax.lax.PrecisionLike#

str | Precision | tuple[str, str] | tuple[Precision, Precision] | None 的别名

class jax.lax.RoundingMethod(value)[source]#

枚举。

class jax.lax.ScatterDimensionNumbers(update_window_dims, inserted_window_dims, scatter_dims_to_operand_dims)[source]#

描述 XLA 的 Scatter 运算符 的维度编号参数。有关维度编号含义的更多详细信息,请参阅 XLA 文档。

参数:
  • update_window_dims (Sequence[int]) – updates 中的一组窗口维度。必须是按升序排列的整数元组,每个整数表示维度编号。

  • inserted_window_dims (Sequence[int]) – 必须插入到 updates 形状中的大小为 1 的窗口维度集。必须是按升序排列的整数元组,每个整数表示输出的维度编号。这些是 gather 情况下 collapsed_slice_dims 的镜像。

  • scatter_dims_to_operand_dims (Sequence[int]) – 对于 scatter_indices 中的每个维度,给出 operand 中的对应维度。必须是大小等于 scatter_indices.shape[-1] 的整数序列。

与 XLA 的 ScatterDimensionNumbers 结构不同,index_vector_dim 是隐式的;始终存在索引向量维度,并且它必须始终是最后一个维度。要散布标量索引,请添加一个大小为 1 的尾随维度。