公共API: jax 包

子包

• jax.numpy 模块
• jax.scipy 模块
• jax.lax 模块
• jax.random 模块
• jax.sharding 模块
• jax.debug 模块
• jax.dlpack 模块
• jax.distributed 模块
• jax.dtypes 模块
• jax.flatten_util 模块
• jax.image 模块
• jax.nn 模块
• jax.ops 模块
• jax.profiler 模块
• jax.stages 模块
• jax.tree 模块
• jax.tree_util 模块
• jax.typing 模块
• jax.export 模块
• jax.extend 模块
• jax.example_libraries 模块
• jax.experimental 模块

配置

config
check_tracer_leaks	用于 jax_check_tracer_leaks 配置选项的上下文管理器。
checking_leaks	用于 jax_check_tracer_leaks 配置选项的上下文管理器。
debug_nans	用于 jax_debug_nans 配置选项的上下文管理器。
debug_infs	用于 jax_debug_infs 配置选项的上下文管理器。
default_device	jax_default_device 配置选项的上下文管理器。
default_matmul_precision	用于 jax_default_matmul_precision 配置选项的上下文管理器。
default_prng_impl	jax_default_prng_impl 配置选项的上下文管理器。
enable_checks	jax_enable_checks 配置选项的上下文管理器。
enable_custom_prng	用于 jax_enable_custom_prng 配置选项的上下文管理器（临时）。
enable_custom_vjp_by_custom_transpose	用于 jax_enable_custom_vjp_by_custom_transpose 配置选项的上下文管理器（瞬态）。
log_compiles	jax_log_compiles 配置选项的上下文管理器。
numpy_rank_promotion	jax_numpy_rank_promotion 配置选项的上下文管理器。
transfer_guard(new_val)	一个上下文管理器，用于控制所有传输的传输保护级别。

即时编译 (jit)

jit(fun[, in_shardings, out_shardings, ...])	为 fun 设置即时编译（JIT）与 XLA。
disable_jit([disable])	在其动态上下文中禁用 jit() 行为的上下文管理器。
ensure_compile_time_eval()	上下文管理器以确保在跟踪/编译时进行评估（或错误）。
xla_computation(fun[, static_argnums, ...])	创建一个函数，该函数在给定示例参数的情况下生成其XLA计算。
make_jaxpr([axis_env, return_shape, ...])	创建一个函数，该函数在给定示例参数的情况下生成其 jaxpr。
eval_shape(fun, *args, **kwargs)	在不进行任何浮点运算的情况下计算 fun 的形状/数据类型。
ShapeDtypeStruct(shape, dtype[, ...])	用于存储数组的形状、数据类型和其他静态属性的容器。
device_put(x[, device, src])	将 x 转移到 device。
device_put_replicated(x, devices)	将数组传输到每个指定的设备并形成数组。
device_put_sharded(shards, devices)	将数组分片传输到指定设备并形成数组。
device_get(x)	将 x 传输到主机。
default_backend()	返回默认XLA后端的平台名称。
named_call(fun, *[, name])	在分阶段执行JAX计算时，为函数添加用户指定的名称。
named_scope(name)	一个上下文管理器，将用户指定的名称添加到 JAX 名称堆栈中。
block_until_ready(x)	尝试在 pytree 叶子上调用 block_until_ready 方法。

自动微分

grad(fun[, argnums, has_aux, holomorphic, ...])	创建一个评估 fun 梯度的函数。
value_and_grad(fun[, argnums, has_aux, ...])	创建一个函数，该函数计算 fun 及其梯度。
jacfwd(fun[, argnums, has_aux, holomorphic])	fun 的雅可比矩阵，通过前向模式自动微分逐列计算。
jacrev(fun[, argnums, has_aux, holomorphic, ...])	fun 的雅可比矩阵按行计算，使用反向模式自动微分。
hessian(fun[, argnums, has_aux, holomorphic])	fun 的 Hessian 矩阵，以密集数组形式表示。
jvp(fun, primals, tangents[, has_aux])	计算 fun 的（前向模式）雅可比向量积。
linearize()	使用 jvp() 和部分求值生成 fun 的线性近似。
linear_transpose(fun, *primals[, reduce_axes])	转置一个保证为线性的函数。
vjp() ))	计算 fun 的（反向模式）向量-雅可比积。
custom_gradient(fun)	用于定义自定义 VJP 规则（即自定义梯度）的便捷函数。
closure_convert(fun, *example_args)	闭包转换工具，用于高阶自定义导数。
checkpoint(fun, *[, prevent_cse, policy, ...])	当微分时，使 fun 重新计算内部的线性化点。

custom_jvp

custom_jvp(fun[, nondiff_argnums])	为自定义 JVP 规则定义设置一个可 JAX 变换的函数。
custom_jvp.defjvp(jvp[, symbolic_zeros])	为此实例表示的函数定义一个自定义的 JVP 规则。
custom_jvp.defjvps(*jvps)	为每个参数分别定义JVP的便捷包装器。

custom_vjp

custom_vjp(fun[, nondiff_argnums])	为自定义 VJP 规则定义设置一个可 JAX 变换的函数。
custom_vjp.defvjp(fwd, bwd[, ...])	为此实例表示的函数定义一个自定义的 VJP 规则。

jax.Array (jax.Array)

Array()	JAX 的数组基类
make_array_from_callback(shape, sharding, ...)	通过从 data_callback 获取的数据返回一个 jax.Array。
make_array_from_single_device_arrays(shape, ...)	从一个设备上的 jax.Array 序列返回一个 jax.Array。
make_array_from_process_local_data(sharding, ...)	使用进程中可用的数据创建分布式张量。

数组属性和方法

Array.addressable_shards	可寻址分片列表。
Array.all([axis, out, keepdims, where])	测试沿给定轴的所有数组元素是否评估为 True。
Array.any([axis, out, keepdims, where])	测试沿给定轴的任何数组元素是否评估为 True。
Array.argmax([axis, out, keepdims])	返回最大值的索引。
Array.argmin([axis, out, keepdims])	返回最小值的索引。
Array.argpartition(kth[, axis])	返回部分排序数组的索引。
Array.argsort([axis, kind, order, stable, ...])	返回对数组进行排序的索引。
Array.astype(dtype[, copy, device])	复制数组并转换为指定的数据类型。
Array.at	用于索引更新功能的辅助属性。
Array.choose(choices[, out, mode])	从多个数组的元素中构建一个数组。
Array.clip([min, max])	返回一个数组，其值被限制在指定范围内。
Array.compress(condition[, axis, out, size, ...])	返回沿给定轴的此数组的选择切片。
Array.conj()	返回数组的复共轭。
Array.conjugate()	返回数组的复共轭。
Array.copy()	返回数组的副本。
Array.copy_to_host_async()	异步地将 Array 复制到主机。
Array.cumprod([axis, dtype, out])	返回数组的累积乘积。
Array.cumsum([axis, dtype, out])	返回数组的累计和。
Array.device	与数组API兼容的设备属性。
Array.diagonal([offset, axis1, axis2])	返回数组中指定的对角线。
Array.dot(b, *[, precision, ...])	计算两个数组的点积。
Array.dtype	数组的 数据类型 (numpy.dtype)。
Array.flat	使用 flatten() 代替。
Array.flatten([order])	将数组展平为1维形状。
Array.global_shards	全局分片列表。
Array.imag	返回数组的虚部。
Array.is_fully_addressable	这个数组是完全可寻址的吗？
Array.is_fully_replicated	这个数组是完全复制的吗？
Array.item(*args)	将数组的一个元素复制到一个标准的 Python 标量并返回它。
Array.itemsize	一个数组元素的字节长度。
Array.max([axis, out, keepdims, initial, where])	返回沿给定轴的数组元素的最大值。
Array.mean([axis, dtype, out, keepdims, where])	返回沿给定轴的数组元素的平均值。
Array.min([axis, out, keepdims, initial, where])	返回沿给定轴的数组元素的最小值。
Array.nbytes	数组元素消耗的总字节数。
Array.ndim	数组的维度数量。
Array.nonzero(*[, fill_value, size])	返回数组中非零元素的索引。
Array.prod([axis, dtype, out, keepdims, ...])	返回数组元素在给定轴上的乘积。
Array.ptp([axis, out, keepdims])	返回给定轴上的峰峰值范围。
Array.ravel([order])	将数组展平为1维形状。
Array.real	返回数组的实部。
Array.repeat(repeats[, axis, ...])	从重复元素构建数组。
Array.reshape(*args[, order])	返回一个包含相同数据但形状不同的新数组。
Array.round([decimals, out])	将数组元素四舍五入到给定的小数位。
Array.searchsorted(v[, side, sorter, method])	在一个已排序的数组中执行二分查找。
Array.shape	数组的形状。
Array.sharding	数组的切片。
Array.size	数组中元素的总数。
Array.sort([axis, kind, order, stable, ...])	返回数组的排序副本。
Array.squeeze([axis])	从数组中移除一个或多个长度为1的轴。
Array.std([axis, dtype, out, ddof, ...])	计算沿给定轴的标准差。
Array.sum([axis, dtype, out, keepdims, ...])	数组元素在给定轴上的总和。
Array.swapaxes(axis1, axis2)	交换数组的两个轴。
Array.take(indices[, axis, out, mode, ...])	从数组中提取元素。
Array.to_device(device, *[, stream])	返回指定设备上的数组副本
Array.trace([offset, axis1, axis2, dtype, out])	返回对角线上的元素之和。
Array.transpose(*args)	返回数组的一个转置副本。
Array.var([axis, dtype, out, ddof, ...])	计算沿指定轴的方差。
Array.view([dtype, type])	返回数组的按位复制，视作新的数据类型。
Array.T	计算全轴数组转置。
Array.mT	计算（批量）矩阵转置。

向量化 (vmap)

vmap(fun[, in_axes, out_axes, axis_name, ...])	向量化映射。
numpy.vectorize(pyfunc, *[, excluded, signature])	定义一个带有广播功能的矢量化函数。

并行化 (pmap)

pmap(fun[, axis_name, in_axes, out_axes, ...])	支持集体操作的并行映射。
devices([backend])	返回给定后端的所有设备列表。
local_devices([process_index, backend, host_id])	类似于 jax.devices()，但仅返回给定进程本地的设备。
process_index([backend])	返回此进程的整数进程索引。
device_count([backend])	返回设备总数。
local_device_count([backend])	返回此进程可寻址的设备数量。
process_count([backend])	返回与后端关联的JAX进程数。
process_indices([backend])	返回与后端关联的所有 JAX 进程索引的列表。

回调

pure_callback(callback, result_shape_dtypes, ...)	调用一个纯Python回调函数。
experimental.io_callback(callback, ...[, ...])	调用一个不纯的Python回调函数。
debug.callback(callback, *args[, ordered])	调用一个可分阶段的 Python 回调。
debug.print(fmt, *args[, ordered])	打印值并在分阶段输出的 JAX 函数中工作。

杂项

Device	一个可用设备的描述。
print_environment_info([return_string])	返回一个包含本地环境及 JAX 安装信息字符串。
live_arrays([platform])	返回 platform 后端中的所有活动数组。
clear_caches()	清除所有编译和暂存缓存。