JAX 中文文档（十四）

dctn(x[, type, s, axes, norm])

计算输入的多维离散余弦变换

idct(x[, type, n, axis, norm])

计算输入的离散余弦变换的逆变换

cho_factor(a[, lower, overwrite_a, check_finite])

基于 Cholesky 的线性求解因式分解

cho_solve(c_and_lower, b[, overwrite_b, …])

使用 Cholesky 分解解线性系统

cholesky(a[, lower, overwrite_a, check_finite])

计算矩阵的 Cholesky 分解。

det(a[, overwrite_a, check_finite])

计算矩阵的行列式

eigh()

计算 Hermitian 矩阵的特征值和特征向量

eigh_tridiagonal(d, e, *[, eigvals_only, …])

解对称实三对角矩阵的特征值问题

expm(A, *[, upper_triangular, max_squarings])

计算矩阵指数

expm_frechet()

计算矩阵指数的 Frechet 导数

funm(A, func[, disp])

评估矩阵值函数

hessenberg()

计算矩阵的 Hessenberg 形式

hilbert(n)

创建阶数为 n 的 Hilbert 矩阵。

inv(a[, overwrite_a, check_finite])

返回方阵的逆矩阵

lu()

计算 LU 分解

lu_factor(a[, overwrite_a, check_finite])

基于 LU 的线性求解因式分解

lu_solve(lu_and_piv, b[, trans, …])

使用 LU 分解解线性系统

polar(a[, side, method, eps, max_iterations])

计算极分解

qr()

计算数组的 QR 分解

rsf2csf(T, Z[, check_finite])

将实数舒尔形式转换为复数舒尔形式。

schur(a[, output])

计算舒尔分解

solve(a, b[, lower, overwrite_a, …])

解线性方程组

solve_triangular(a, b[, trans, lower, …])

解上（或下）三角线性方程组

sqrtm(A[, blocksize])

计算矩阵的平方根

svd()

计算奇异值分解

convolve(in1, in2[, mode, method, precision])

两个 N 维数组的卷积。

convolve2d(in1, in2[, mode, boundary, …])

两个二维数组的卷积。

correlate(in1, in2[, mode, method, precision])

两个 N 维数组的互相关。

correlate2d(in1, in2[, mode, boundary, …])

两个二维数组的互相关。

csd(x, y[, fs, window, nperseg, noverlap, …])

使用 Welch 方法估计交叉功率谱密度（CSD）。

detrend(data[, axis, type, bp, overwrite_data])

从数据中移除线性或分段线性趋势。

istft(Zxx[, fs, window, nperseg, noverlap, …])

执行逆短时傅里叶变换（ISTFT）。

stft(x[, fs, window, nperseg, noverlap, …])

计算短时傅里叶变换（STFT）。

cg(A, b[, x0, tol, atol, maxiter, M])

使用共轭梯度法解决 Ax = b。

beta()

贝塔函数

betainc(a, b, x)

正则化的不完全贝塔函数。

betaln(a, b)

贝塔函数绝对值的自然对数

digamma(x)

Digamma 函数

entr(x)

熵函数

erf(x)

误差函数

erfc(x)

误差函数的补函数

erfinv(x)

误差函数的反函数

exp1(x)

指数积分函数。

expi

指数积分函数。

expit(x)

逻辑 sigmoid（expit）函数

expn

广义指数积分函数。

factorial(n[, exact])

阶乘函数

gamma(x)

伽马函数。

gammainc(a, x)

正则化的下不完全伽马函数。

gammaincc(a, x)

正则化的上不完全伽马函数。

gammaln(x)

伽马函数绝对值的自然对数。

gammasgn(x)

伽马函数的符号。

hyp1f1

1F1 超几何函数。

i0(x)

修改贝塞尔函数零阶。

i0e(x)

指数缩放的修改贝塞尔函数零阶。

i1(x)

修改贝塞尔函数一阶。

i1e(x)

指数缩放的修改贝塞尔函数一阶。

log_ndtr

对数正态分布函数。

logit

对数几率函数。

logsumexp()

对数-总和-指数归约。

lpmn(m, n, z)

第一类相关勒让德函数（ALFs）。

lpmn_values(m, n, z, is_normalized)

第一类相关勒让德函数（ALFs）。

multigammaln(a, d)

多变量伽马函数的自然对数。

ndtr(x)

正态分布函数。

ndtri§

正态分布函数的反函数。

poch

Pochhammer 符号。

polygamma(n, x)

多次伽马函数。

spence(x)

斯宾斯函数，也称实数域下的二元对数函数。

sph_harm(m, n, theta, phi[, n_max])

计算球谐函数。

xlog1py

计算 x*log(1 + y)，当 x=0 时返回 0。

xlogy

计算 x*log(y)，当 x=0 时返回 0。

zeta

赫维茨 ζ 函数。

kl_div(p, q)

库尔巴克-莱布勒散度。

rankdata(a[, method, axis, nan_policy])

计算数组沿轴的排名。

sem(a[, axis, ddof, nan_policy, keepdims])

计算均值的标准误差。

pmf(k, p[, loc])

伯努利概率质量函数。

cdf(k, p)

伯努利累积分布函数。

pdf(x, a, b[, loc, scale])

Beta 概率分布函数。

cdf(x, a, b[, loc, scale])

Beta 累积分布函数。

logcdf(x, a, b[, loc, scale])

Beta 对数累积分布函数。

sf(x, a, b[, loc, scale])

Beta 分布生存函数。

pdf(x[, loc, scale])

柯西概率分布函数。

cdf(x[, loc, scale])

柯西累积分布函数。

logcdf(x[, loc, scale])

柯西对数累积分布函数。

sf(x[, loc, scale])

柯西分布对数生存函数。

logsf(x[, loc, scale])

柯西对数生存函数。

isf(q[, loc, scale])

柯西分布逆生存函数。

pdf(x, df[, loc, scale])

卡方概率分布函数。

cdf(x, df[, loc, scale])

卡方累积分布函数。

logcdf(x, df[, loc, scale])

卡方对数累积分布函数。

sf(x, df[, loc, scale])

卡方生存函数。

pdf(x, a[, loc, scale])

伽玛概率分布函数。

cdf(x, a[, loc, scale])

伽玛累积分布函数。

logcdf(x, a[, loc, scale])

伽玛对数累积分布函数。

sf(x, a[, loc, scale])

伽玛生存函数。

logpdf(x, beta)

广义正态对数概率分布函数。

logpdf(x[, loc, scale])

拉普拉斯对数概率分布函数。

isf(x[, loc, scale])

Logistic 分布逆生存函数。

logpdf(x[, loc, scale])

Logistic 对数概率分布函数。

pdf(x[, loc, scale])

Logistic 概率分布函数。

ppf(x[, loc, scale])

Logistic 分位点函数。

pdf(x[, loc, scale])

正态分布概率分布函数。

cdf(x[, loc, scale])

正态分布累积分布函数。

logcdf(x[, loc, scale])

正态分布对数累积分布函数。

ppf(q[, loc, scale])

正态分布百分点函数。

sf(x[, loc, scale])

正态分布生存函数。

logsf(x[, loc, scale])

正态分布对数生存函数。

pmf(k, mu[, loc])

泊松分布概率质量函数。

logcdf(x, a, b[, loc, scale])

截断正态分布对数累积分布函数。

logpdf(x, a, b[, loc, scale])

截断正态分布对数概率分布函数。

logsf(x, a, b[, loc, scale])

截断正态分布对数生存函数。

pdf(x, a, b[, loc, scale])

截断正态分布概率分布函数。

pdf(x[, loc, scale])

均匀分布概率分布函数。

cdf(x[, loc, scale])

均匀分布累积分布函数。

ppf(q[, loc, scale])

均匀分布百分点函数。

gaussian_kde.evaluate(points)

对给定点评估高斯核密度估计器。

gaussian_kde.integrate_gaussian(mean, cov)

加权高斯积分分布。

gaussian_kde.integrate_box_1d(low, high)

在给定限制下积分分布。

gaussian_kde.integrate_kde(other)

集成两个高斯核密度估计分布的乘积。

gaussian_kde.resample(key[, shape])

从估计的概率密度函数中随机采样数据集

gaussian_kde.pdf(x)

概率密度函数

gaussian_kde.logpdf(x)

对数概率密度函数

pdf(x, c)

Wrapped Cauchy 概率分布函数。

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)

按元素计算 popcount，即每个元素中设置的位数。

broadcast(operand, sizes)

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

broadcast_in_dim(operand, shape, …)

包装 XLA 的 BroadcastInDim 操作符。

broadcast_shapes()

返回经过 NumPy 广播后的形状。

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)

元素级 clamp 函数。

clz(x)

元素级计算前导零的个数。

collapse(operand, start_dimension[, …])

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

complex(x, y)

元素级构造复数：(x + jy)。

concatenate(operands, 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, …)

将卷积维度编号转换为 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])

沿轴计算累积 logsumexp。

cummax(operand[, axis, reverse])

沿轴计算累积最大值。

cummin(operand[, axis, reverse])

沿轴计算累积最小值。

cumprod(operand[, axis, reverse])

沿轴计算累积乘积。

cumsum(operand[, axis, reverse])

沿轴计算累积和。

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])

返回填充值为 fill_value 的形状数组。

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

基于示例数组 x 创建类似于 np.full 的完整数组。

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

Gather 操作符。

ge(x, y)

元素级大于或等于：(x \geq y)。

gt(x, y)

元素级大于：(x > y)。

igamma(a, x)

元素级正则化不完全 gamma 函数。

igammac(a, x)

元素级补充正则化不完全 gamma 函数。

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)

元素级对数 gamma 函数：(\mathrm{log}(\Gamma(x)))。

log(x)

元素级自然对数：(\mathrm{log}(x))。

log1p(x)

元素级 (\mathrm{log}(1 + x))。

logistic(x)

元素级 logistic（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)

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

pad(operand, padding_value, padding_config)

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

polygamma(m, x)

按位多次 gamma 函数：(\psi^{(m)}(x))。

population_count(x)

按位人口统计，统计每个元素中设置的位数。

pow(x, y)

按位幂运算：(x^y)。

random_gamma_grad(a, x)

Gamma 分布导数的按位计算。

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])

无状态的伪随机数位生成器。

rng_uniform(a, b, shape)

有状态的伪随机数生成器。

round(x[, rounding_method])

按位四舍五入。

rsqrt(x)

按位倒数平方根：(1 \over \sqrt{x})。

scatter(operand, scatter_indices, updates, …)

Scatter-update 运算符。

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

Scatter-add 运算符。

scatter_apply(operand, scatter_indices, …)

Scatter-apply 运算符。

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

Scatter-max 运算符。

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

Scatter-min 运算符。

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

Scatter-multiply 运算符。

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() 的单维度应用封装。

sort()

封装了 XLA 的 Sort 运算符。

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

沿着dimension排序keys并对values应用相同的置换。

sqrt(x)

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

square(x)

逐元素平方：(x²)。

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))

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

根据条件应用true_fun或false_fun。

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

通过归约到jax.lax.while_loop()从lower到upper循环。

map(f, xs)

在主要数组轴上映射函数。

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。

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

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

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

可微分求解函数的根。

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

映射轴的实例化和映射不同轴。

pdot(x, y, axis_name[, pos_contract, …])

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

在映射的轴axis_name上进行全归约求和。

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

像psum(x, axis_name)，但每个设备仅保留部分结果。

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

在映射的轴axis_name上计算全归约最大值。

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

在映射的轴axis_name上计算全归约最小值。

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

在映射的轴axis_name上计算全归约均值。

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 的索引。

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

一般矩阵的特征分解。

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

Hermite 矩阵的特征分解。

hessenberg(a)

将方阵约化为上 Hessenberg 形式。

lu(x)

带有部分主元列主元分解。

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])

将对称/Hermitian 矩阵约化为三对角形式。

tridiagonal_solve(dl, d, du, b)

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

在 TPU 上最快

✅

✅

✅

✅

可以高效分片（使用 pjit）

✅

✅

✅

在分片中相同

✅

✅

✅

✅

在 CPU/GPU/TPU 上相同

✅

✅

在 JAX/XLA 版本间相同

✅

✅

key(seed, *[, impl])

给定整数种子创建伪随机数生成器（PRNG）密钥。

key_data（密钥）

恢复 PRNG 密钥数组下的密钥数据位。

wrap_key_data(key_bits_array, *[, impl])

将密钥数据位数组包装成 PRNG 密钥数组。

fold_in(key, data)

将数据折叠到 PRNG 密钥中，形成新的 PRNG 密钥。

split(key[, num])

将 PRNG 密钥按添加一个前导轴拆分为 num 个新密钥。

clone(key)

克隆一个密钥以便重复使用。

bernoulli(key[, p, shape])

采样给定形状和均值的伯努利分布随机值。

beta(key, a, b[, shape, dtype])

采样给定形状和浮点数数据类型的贝塔分布随机值。

binomial(key, n, p[, shape, dtype])

采样给定形状和浮点数数据类型的二项分布随机值。

bits(key[, shape, dtype])

以无符号整数的形式采样均匀比特。

categorical(key, logits[, axis, shape])

从分类分布中采样随机值。

cauchy(key[, shape, dtype])

采样给定形状和浮点数数据类型的柯西分布随机值。

chisquare(key, df[, shape, dtype])

采样给定形状和浮点数数据类型的卡方分布随机值。

choice(key, a[, shape, replace, p, axis])

从给定数组中生成随机样本。

dirichlet(key, alpha[, shape, dtype])

采样给定形状和浮点数数据类型的狄利克雷分布随机值。

double_sided_maxwell(key, loc, scale[, …])

从双边 Maxwell 分布中采样。

exponential(key[, shape, dtype])

采样给定形状和浮点数数据类型的指数分布随机值。

f(key, dfnum, dfden[, shape, dtype])

采样给定形状和浮点数数据类型的 F 分布随机值。

gamma(key, a[, shape, dtype])

采样给定形状和浮点数数据类型的伽马分布随机值。

generalized_normal(key, p[, shape, dtype])

从广义正态分布中采样。

geometric(key, p[, shape, dtype])

采样给定形状和浮点数数据类型的几何分布随机值。

gumbel(key[, shape, dtype])

采样给定形状和浮点数数据类型的 Gumbel 分布随机值。

laplace(key[, shape, dtype])

采样给定形状和浮点数数据类型的拉普拉斯分布随机值。

loggamma(key, a[, shape, dtype])

采样给定形状和浮点数数据类型的对数伽马分布随机值。

logistic(key[, shape, dtype])

采样给定形状和浮点数数据类型的 logistic 随机值。

lognormal(key[, sigma, shape, dtype])

采样给定形状和浮点数数据类型的对数正态分布随机值。

maxwell(key[, shape, dtype])

从单边 Maxwell 分布中采样。

multivariate_normal(key, mean, cov[, shape, …])

采样给定均值和协方差的多变量正态分布随机值。

normal(key[, shape, dtype])

采样给定形状和浮点数数据类型的标准正态分布随机值。

orthogonal(key, n[, shape, dtype])

从正交群 O(n) 中均匀采样。

pareto(key, b[, shape, dtype])

采样给定形状和浮点数数据类型的帕累托分布随机值。

permutation(key, x[, axis, independent])

返回随机排列的数组或范围。

poisson(key, lam[, shape, dtype])

采样给定形状和整数数据类型的泊松分布随机值。

rademacher(key[, shape, dtype])

从 Rademacher 分布中采样。

randint(key, shape, minval, maxval[, dtype])

用给定的形状和数据类型在[minval, maxval)范围内示例均匀随机整数值。

[rayleigh(key, scale[, shape, dtype])

用给定的形状和浮点数数据类型示例瑞利随机值。

t(key, df[, shape, dtype])

用给定的形状和浮点数数据类型示例学生 t 分布随机值。

triangular(key, left, mode, right[, shape, …])

用给定的形状和浮点数数据类型示例三角形随机值。

truncated_normal(key, lower, upper[, shape, …])

用给定的形状和数据类型示例截断标准正态随机值。

uniform(key[, shape, dtype, minval, maxval])

用给定的形状和数据类型在[minval, maxval)范围内示例均匀随机值。

[wald(key, mean[, shape, dtype])

用给定的形状和浮点数数据类型示例瓦尔德随机值。

weibull_min(key, scale, concentration[, …])

从威布尔分布中采样。

print(fmt, *args[, ordered])

打印值，并在 JAX 函数中工作。

breakpoint(*[, backend, filter_frames, …])

在程序中某一点设置断点。

visualize_array_sharding(arr, **kwargs)

可视化数组的分片。

visualize_sharding(shape, sharding, *[, …])

使用 rich 可视化 Sharding。

to_dlpack(x[, stream, src_device, …])

返回一个封装了 Array x 的 DLPack 张量。

shutdown()

关闭分布式系统。

canonicalize_dtype(dtype[, allow_extended_dtype])

根据config.x64_enabled配置将 dtype 转换为规范的 dtype。

float0

对应于相同名称的标量类型和 dtype 的 DType 类。

issubdtype(a, b)

如果第一个参数是类型代码在类型层次结构中较低/相等，则返回 True。

prng_key()

PRNG Key dtypes 的标量类。

result_type(*args[, return_weak_type_flag])

方便函数，用于应用 JAX 参数 dtype 提升。

scalar_type_of(x)

返回与 JAX 值关联的标量类型。

scale_and_translate(image, shape, …[, …])

对图像应用缩放和平移。

relu6

线性整流单元 6 激活函数。

sigmoid(x)

Sigmoid 激活函数。

softplus(x)

Softplus 激活函数。

sparse_plus(x)

稀疏加法函数。

sparse_sigmoid(x)

稀疏 Sigmoid 激活函数。

soft_sign(x)

Soft-sign 激活函数。

silu(x)

SiLU（又称 swish）激活函数。

swish(x)

SiLU（又称 swish）激活函数。

log_sigmoid(x)

对数 Sigmoid 激活函数。

leaky_relu(x[, negative_slope])

泄漏整流线性单元激活函数。

hard_sigmoid(x)

硬 Sigmoid 激活函数。

hard_silu(x)

硬 SiLU（swish）激活函数。

hard_swish(x)

硬 SiLU（swish）激活函数。

hard_tanh(x)

硬\tanh 激活函数。

elu(x[, alpha])

指数线性单元激活函数。

celu(x[, alpha])

连续可微的指数线性单元激活函数。

selu(x)

缩放的指数线性单元激活函数。

gelu(x[, approximate])

高斯误差线性单元激活函数。

glu(x[, axis])

门控线性单元激活函数。

squareplus(x[, b])

Squareplus 激活函数。

mish(x)

Mish 激活函数。

log_softmax(x[, axis, where, initial])

对数 Softmax 函数。

logsumexp()

对数-总和-指数归约。

standardize(x[, axis, mean, variance, …])

通过减去mean并除以(\sqrt{\mathrm{variance}})来标准化数组。

one_hot(x, num_classes, *[, dtype, axis])

对给定索引进行 One-hot 编码。

delta_orthogonal([scale, column_axis, dtype])

构建一个用于增量正交核的初始化器。

glorot_normal([in_axis, out_axis, …])

构建一个 Glorot 正态初始化器（又称 Xavier 正态初始化器）。

glorot_uniform([in_axis, out_axis, …])

构建一个 Glorot 均匀初始化器（又称 Xavier 均匀初始化器）。

he_normal([in_axis, out_axis, batch_axis, dtype])

构建一个 He 正态初始化器（又称 Kaiming 正态初始化器）。

he_uniform([in_axis, out_axis, batch_axis, …])

构建一个 He 均匀初始化器（又称 Kaiming 均匀初始化器）。

lecun_normal([in_axis, out_axis, …])

构建一个 Lecun 正态初始化器。

lecun_uniform([in_axis, out_axis, …])

构建一个 Lecun 均匀初始化器。

normal([stddev, dtype])

构建一个返回实数正态分布随机数组的初始化器。

ones(key, shape[, dtype])

返回一个填充为一的常数数组的初始化器。

orthogonal([scale, column_axis, dtype])

构建一个返回均匀分布正交矩阵的初始化器。

truncated_normal([stddev, dtype, lower, upper])

构建一个返回截断正态分布随机数组的初始化器。

uniform([scale, dtype])

构建一个返回实数均匀分布随机数组的初始化器。

variance_scaling(scale, mode, distribution)

初始化器，根据权重张量的形状调整其尺度。

zeros(key, shape[, dtype])

返回一个填充零的常数数组的初始化器。

segment_min(data, segment_ids[, …])

segment_prod(data, segment_ids[, …])

segment_sum(data, segment_ids[, …])

start_trace(log_dir[, create_perfetto_link, …])

启动性能分析跟踪。

stop_trace()

停止当前正在运行的性能分析跟踪。

trace(log_dir[, create_perfetto_link, …])

上下文管理器，用于进行性能分析跟踪。

annotate_function(func[, name])

生成函数执行的跟踪事件的装饰器。

TraceAnnotation

在分析器中生成跟踪事件的上下文管理器。

StepTraceAnnotation(name, **kwargs)

在分析器中生成步骤跟踪事件的上下文管理器。

save_device_memory_profile(filename[, backend])

收集设备内存使用情况，并将其写入文件。