公共API：jax 包#

子包#

配置#

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配置选项的上下文管理器（瞬时）。

enable_x64

用于jax_enable_x64配置选项的上下文管理器。

log_compiles

用于jax_log_compiles配置选项的上下文管理器。

no_tracing

用于jax_no_tracing配置选项的上下文管理器。

numpy_rank_promotion

用于jax_numpy_rank_promotion配置选项的上下文管理器。

transfer_guard(new_val)

一个上下文管理器，用于控制所有传输的传输保护级别。

即时编译（jit）#

jit(fun, /, *[, in_shardings, ...])

为XLA的即时编译设置fun。

disable_jit([disable])

禁用其动态上下文下jit()行为的上下文管理器。

ensure_compile_time_eval()

确保在跟踪/编译时进行评估（或出错）的上下文管理器。

make_jaxpr([axis_env, return_shape, ...])

给定示例参数，创建一个返回fun的jaxpr的函数。

eval_shape(fun, *args, **kwargs)

在不进行任何FLOPs的情况下计算fun的形状/dtype。

ShapeDtypeStruct(shape, dtype, *[, ...])

一个用于存储数组的形状、dtype和其他静态属性的容器。

device_put(x[, device, src, donate, may_alias])

将x传输到device。

device_get(x)

将x传输到主机。

default_backend()

返回默认XLA后端平台的名称。

named_call(fun, *[, name])

在暂存JAX计算时，向函数添加用户指定的名称。

named_scope(name)

一个上下文管理器，将用户指定的名称添加到JAX名称堆栈中。

block_until_ready(x)

尝试调用pytree叶子上的block_until_ready方法。

copy_to_host_async(x)

尝试调用pytree叶子上的copy_to_host_async方法。

make_mesh(axis_shapes, axis_names[, ...])

使用指定的形状和轴名称创建高效的网格。

自动微分#

grad(fun[, argnums, has_aux, holomorphic, ...])

创建一个评估fun梯度的函数。

value_and_grad(fun[, argnums, has_aux, ...])

创建一个同时评估fun及其梯度的函数。

jacobian(fun[, argnums, has_aux, ...])

别名为jax.jacrev()。

jacfwd(fun[, argnums, has_aux, holomorphic])

使用前向模式AD逐列计算的fun的雅可比矩阵。

jacrev(fun[, argnums, has_aux, holomorphic, ...])

使用反向模式AD逐行计算的fun的雅可比矩阵。

hessian(fun[, argnums, has_aux, holomorphic])

作为密集数组的fun的海森矩阵。

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在微分时重新计算内部线性化点。

矢量化#

vmap(fun[, in_axes, out_axes, axis_name, ...])

矢量化映射。

numpy.vectorize(pyfunc, *[, excluded, signature])

定义一个支持广播的向量化函数。

并行化#

shard_map([f, in_specs, mesh, axis_names, ...])

使用设备网格对数据分片进行映射。

smap([f, in_axes])

单轴分片映射，一次映射一个轴的函数f。

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进程索引列表。

定制#

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 规则。

custom_batching#

custom_batching.custom_vmap(fun)

自定义JAX可转换函数的vmap行为。

custom_batching.custom_vmap.def_vmap(vmap_rule)

为此 `custom_vmap` 函数定义 `vmap` 规则。

custom_batching.sequential_vmap(f)

使用循环的custom_vmap的特殊情况。

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.committed

数组是否已提交。

Array.conj()

返回数组的复共轭。

Array.conjugate()

返回数组的复共轭。

Array.copy()

返回数组的副本。

Array.copy_to_host_async()

将 Array 异步复制到主机。

Array.cumprod([axis, dtype, out])

返回数组的累积乘积。

Array.cumsum([axis, dtype, out])

返回数组的累积和。

Array.device

兼容 Array API 的设备属性。

Array.diagonal([offset, axis1, axis2])

从数组中返回指定的对角线。

Array.dot(b, *[, precision, ...])

计算两个数组的点积。

Array.dtype

数组的数据类型（numpy.dtype）。

Array.flat

请使用flatten()代替。

Array.flatten([order, out_sharding])

将数组展平为一维形状。

Array.global_shards

全局分片列表。

Array.imag

返回数组的虚部。

Array.is_fully_addressable

此 Array 是否完全可寻址？

Array.is_fully_replicated

此 Array 是否完全复制？

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

将数组展平为一维形状。

Array.real

返回数组的实部。

Array.repeat(repeats[, axis, ...])

从重复元素构造数组。

Array.reshape(*args[, order, out_sharding])

返回一个包含相同数据但具有新形状的数组。

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

计算（批处理）矩阵转置。

回调#

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

清除所有编译和暂存缓存。

检查点策略#

checkpoint_policies.everything_saveable(**__)

默认策略，如同未使用jax.checkpoint一样。

checkpoint_policies.nothing_saveable(**__)

重新计算所有内容，如同未使用自定义策略一样。

checkpoint_policies.dots_saveable(*_, **__)

checkpoint_policies.checkpoint_dots(*_, **__)

checkpoint_policies.dots_with_no_batch_dims_saveable(...)

这是Transformer的一个有用启发式方法。

checkpoint_policies.checkpoint_dots_with_no_batch_dims(...)

这是Transformer的一个有用启发式方法。

checkpoint_policies.save_any_names_but_these()

仅保存命名值，即checkpoint_name的任何输出，排除给定的名称。

checkpoint_policies.save_only_these_names()

仅保存命名值，并且仅限于给定的名称。

checkpoint_policies.offload_dot_with_no_batch_dims(...)

与dots_with_no_batch_dims_saveable相同，但将数据卸载到CPU内存而不是重新计算。

checkpoint_policies.save_and_offload_only_these_names(*, ...)

与save_only_these_names相同，但将数据卸载到CPU内存而不是重新计算。

checkpoint_policies.save_from_both_policies(...)

给定策略的逻辑OR。