公共 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
模块
配置#
用于 jax_check_tracer_leaks 配置选项的上下文管理器。 |
|
用于 jax_check_tracer_leaks 配置选项的上下文管理器。 |
|
用于 jax_debug_nans 配置选项的上下文管理器。 |
|
用于 jax_debug_infs 配置选项的上下文管理器。 |
|
用于 jax_default_device 配置选项的上下文管理器。 |
|
用于 jax_default_matmul_precision 配置选项的上下文管理器。 |
|
用于 jax_default_prng_impl 配置选项的上下文管理器。 |
|
用于 jax_enable_checks 配置选项的上下文管理器。 |
|
用于 jax_enable_custom_prng 配置选项的上下文管理器 (瞬态)。 |
|
用于 jax_enable_custom_vjp_by_custom_transpose 配置选项的上下文管理器 (瞬态)。 |
|
jax_log_compiles 配置选项的上下文管理器。 |
|
jax_numpy_rank_promotion 配置选项的上下文管理器。 |
|
|
一个用于控制所有传输的传输保护级别的上下文管理器。 |
即时编译 (jit
)#
|
使用 XLA 为即时编译设置 |
|
上下文管理器,在动态上下文中禁用 |
上下文管理器,确保在跟踪/编译时进行评估(或错误)。 |
|
|
创建一个函数,在给定示例参数的情况下生成其 jaxpr。 |
|
计算 |
|
一个包含数组形状、数据类型和其他静态属性的容器。 |
|
将 |
|
将数组(s)传输到每个指定的设备,并形成 Array(s)。 |
|
将数组分片传输到指定的设备,并形成 Array(s)。 |
|
将 |
返回默认 XLA 后端的平台名称。 |
|
|
在将 JAX 计算阶段化时,为函数添加用户指定的名称。 |
|
一个上下文管理器,向 JAX 名称栈添加用户指定的名称。 |
尝试在 pytree 叶子节点上调用 |
自动微分#
|
创建一个函数,评估 |
|
创建一个函数,同时评估 |
|
使用前向模式 AD 列式评估 |
|
使用反向模式 AD 行式评估 |
|
作为密集数组的 |
|
计算 |
使用 |
|
|
转置保证是线性的函数。 |
|
计算 |
|
定义自定义 VJP 规则(也称为自定义梯度)的便利函数。 |
|
闭包转换实用程序,用于与高阶自定义导数一起使用。 |
|
使 |
custom_jvp
#
|
为自定义 JVP 规则定义设置一个 JAX 可转换函数。 |
|
为该实例表示的函数定义一个自定义 JVP 规则。 |
|
为每个参数分别定义 JVP 的便利包装器。 |
custom_vjp
#
|
为自定义 VJP 规则定义设置一个 JAX 可转换函数。 |
|
为该实例表示的函数定义一个自定义 VJP 规则。 |
jax.Array (jax.Array
)#
|
JAX 的 Array 基类 |
|
通过从 |
|
从每个在单个设备上的 |
|
使用进程中可用的数据创建分布式张量。 |
数组属性和方法#
可寻址分片的列表。 |
|
|
测试给定轴上的所有数组元素是否都计算为 True。 |
|
测试给定轴上的任何数组元素是否都计算为 True。 |
|
返回最大值的索引。 |
|
返回最小值的索引。 |
|
返回部分排序数组的索引。 |
|
返回排序数组的索引。 |
|
复制数组并转换为指定的数据类型。 |
索引更新功能的辅助属性。 |
|
|
构造一个从多个数组的元素中选择的数组。 |
|
返回一个值限制在指定范围内的数组。 |
|
沿着给定轴返回此数组的选定切片。 |
返回数组的复共轭。 |
|
返回数组的复共轭。 |
|
返回数组的副本。 |
|
异步将 |
|
|
返回数组的累积积。 |
|
返回数组的累积和。 |
与 Array API 兼容的设备属性。 |
|
|
返回数组中的指定对角线。 |
|
计算两个数组的点积。 |
数组的数据类型 ( |
|
请使用 |
|
|
将数组展平为一维形状。 |
全局分片的列表。 |
|
返回数组的虚部。 |
|
此 Array 是否完全可寻址? |
|
此 Array 是否完全复制? |
|
|
将数组的元素复制到标准 Python 标量并返回它。 |
一个数组元素的字节长度。 |
|
|
返回给定轴上数组元素的最大值。 |
|
返回给定轴上数组元素的平均值。 |
|
返回给定轴上数组元素的最小值。 |
数组元素消耗的总字节数。 |
|
数组的维度数。 |
|
|
返回数组中非零元素的索引。 |
|
返回给定轴上数组元素的乘积。 |
|
返回给定轴上的峰峰值范围。 |
|
将数组展平为一维形状。 |
返回数组的实部。 |
|
|
从重复的元素构造一个数组。 |
|
返回包含相同数据但具有新形状的数组。 |
|
将数组元素舍入到给定的十进制数。 |
|
在排序数组中执行二进制搜索。 |
数组的形状。 |
|
数组的分片。 |
|
数组中的元素总数。 |
|
|
返回数组的排序副本。 |
|
从数组中删除一个或多个长度为 1 的轴。 |
|
计算给定轴上的标准差。 |
|
给定轴上数组元素的总和。 |
|
交换数组的两个轴。 |
|
从数组中获取元素。 |
|
返回指定设备上数组的副本 |
|
返回对角线上的总和。 |
|
返回具有转置轴的数组副本。 |
|
计算给定轴上的方差。 |
|
返回数组的按位副本,视为新的 dtype。 |
计算所有轴数组的转置。 |
|
计算(批处理)矩阵转置。 |
向量化 (vmap
)#
|
向量化映射。 |
|
使用广播定义向量化函数。 |
并行化 (pmap
)#
|
支持集体操作的并行映射。 |
|
返回给定后端的所有设备的列表。 |
|
与 |
|
返回此进程的整数进程索引。 |
|
返回设备总数。 |
|
返回此进程可寻址的设备数量。 |
|
返回与后端关联的 JAX 进程数。 |
|
返回与后端关联的所有 JAX 进程索引的列表。 |
回调#
|
调用纯 Python 回调。 |
|
调用不纯的 Python 回调。 |
|
调用可暂存的 Python 回调。 |
|
打印值并在暂存的 JAX 函数中工作。 |
杂项#
可用设备的描述符。 |
|
|
返回包含本地环境和 JAX 安装信息的字符串。 |
|
返回 platform 后端的后台所有活动数组。 |
清除所有编译和暂存缓存。 |