JAX 中文文档（十四）（5）

JAX 中文文档（十四）（4）https://developer.aliyun.com/article/1559759

jax.debug 模块

原文：jax.readthedocs.io/en/latest/jax.debug.html

运行时值调试实用工具

jax.debug.print 和 jax.debug.breakpoint 描述了如何利用 JAX 的运行时值调试功能。

调试分片实用工具

能够在分段函数内（和外部）检查和可视化数组分片的函数。

jax.dlpack 模块

原文：jax.readthedocs.io/en/latest/jax.dlpack.html

jax.distributed 模块

原文：jax.readthedocs.io/en/latest/jax.distributed.html

jax.dtypes 模块

原文：jax.readthedocs.io/en/latest/jax.dtypes.html

jax.flatten_util 模块

原文：jax.readthedocs.io/en/latest/jax.flatten_util.html

函数列表

jax.image 模块

原文：jax.readthedocs.io/en/latest/jax.image.html

图像操作函数。

更多的图像操作函数可以在建立在 JAX 之上的库中找到，例如 PIX。

图像操作函数

参数类

class jax.image.ResizeMethod(value)

图像调整大小方法。

可能的取值包括：

NEAREST：

最近邻插值。

LINEAR：

线性插值。

LANCZOS3：

Lanczos 重采样，使用半径为 3 的核。

LANCZOS5：

Lanczos 重采样，使用半径为 5 的核。

CUBIC：

三次插值，使用 Keys 三次核。

jax.nn 模块

原文：jax.readthedocs.io/en/latest/jax.nn.html

• jax.nn.initializers 模块

神经网络库常见函数。

激活函数

其他函数

jax.nn.initializers 模块

原文：jax.readthedocs.io/en/latest/jax.nn.initializers.html

与 Keras 和 Sonnet 中定义一致的常见神经网络层初始化器。

初始化器

该模块提供了与 Keras 和 Sonnet 中定义一致的常见神经网络层初始化器。

初始化器是一个函数，接受三个参数：(key, shape, dtype)，并返回一个具有形状shape和数据类型dtype的数组。参数key是一个 PRNG 密钥（例如来自jax.random.key()），用于生成初始化数组的随机数。

jax.ops 模块

原文：jax.readthedocs.io/en/latest/jax.ops.html

段落约简运算符

| segment_max(data, segment_ids[, …]) | 计算数组段内的最大值。 |

jax.profiler 模块

原文：jax.readthedocs.io/en/latest/jax.profiler.html

跟踪和时间分析

描述了如何利用 JAX 的跟踪和时间分析功能进行程序性能分析。

设备内存分析

请参阅设备内存分析，了解 JAX 的设备内存分析功能简介。

jax.stages 模块

原文：jax.readthedocs.io/en/latest/jax.stages.html

接口到编译执行过程的各个阶段。

JAX 转换，例如jax.jit和jax.pmap，也支持一种通用的显式降阶和预编译执行 ahead of time 的方式。 该模块定义了代表这一过程各个阶段的类型。

有关更多信息，请参阅AOT walkthrough。

类

class jax.stages.Wrapped(*args, **kwargs)

一个准备好进行追踪、降阶和编译的函数。

此协议反映了诸如jax.jit之类的函数的输出。 调用它会导致 JIT（即时）降阶、编译和执行。 它也可以在编译之前明确降阶，并在执行之前编译结果。

__call__(*args, **kwargs)

执行包装的函数，根据需要进行降阶和编译。

lower(*args, **kwargs)

明确为给定的参数降阶此函数。

一个降阶函数被从 Python 阶段化，并翻译为编译器的输入语言，可能以依赖于后端的方式。 它已准备好进行编译，但尚未编译。

返回：

一个Lowered实例，表示降阶。

返回类型：

降阶

trace(*args, **kwargs)

明确为给定的参数追踪此函数。

一个追踪函数被从 Python 阶段化，并翻译为一个 jaxpr。 它已准备好进行降阶，但尚未降阶。

返回：

一个Traced实例，表示追踪。

返回类型：

追踪

class jax.stages.Lowered(lowering, args_info, out_tree, no_kwargs=False)

降阶一个根据参数类型和值特化的函数。

降阶是一种准备好进行编译的计算。 此类将降阶与稍后编译和执行所需的剩余信息一起携带。 它还提供了一个通用的 API，用于查询 JAX 各种降阶路径（jit()、pmap()等）中降阶计算的属性。

参数：

• 降阶（XlaLowering）
  
• args_info（Any）
  
• out_tree（PyTreeDef）
  
• no_kwargs（bool）
  

as_text(dialect=None)

此降阶的人类可读文本表示。

旨在可视化和调试目的。 这不必是有效的也不一定可靠的序列化。 它直接传递给外部调用者。

参数：

方言（str | 无) – 可选字符串，指定一个降阶方言（例如，“stablehlo”）

返回类型：

str

compile(compiler_options=None)

编译，并返回相应的Compiled实例。

参数：

compiler_options (dict[str, str | bool] | None)

返回类型：

Compiled

compiler_ir(dialect=None)

这种降低的任意对象表示。

旨在调试目的。这不是有效的也不是可靠的序列化。输出在不同调用之间没有一致性的保证。

如果不可用，则返回None，例如基于后端、编译器或运行时。

参数：

dialect (str | None) – 可选字符串，指定一个降低方言（例如“stablehlo”）

返回类型：

Any | None

cost_analysis()

执行成本估算的摘要。

旨在可视化和调试。此输出的对象是一些简单的数据结构，可以轻松打印或序列化（例如，带有数值叶的嵌套字典、列表和元组）。然而，它的结构可以是任意的：在 JAX 和 jaxlib 的不同版本甚至调用之间可能不一致。

如果不可用，则返回None，例如基于后端、编译器或运行时。

返回类型：

Any | None

property in_tree: PyTreeDef

一对（位置参数、关键字参数）的树结构。

class jax.stages.Compiled(executable, args_info, out_tree, no_kwargs=False)

编译后的函数专门针对类型/值进行了优化表示。

编译计算与可执行文件相关联，并提供执行所需的剩余信息。它还为查询 JAX 的各种编译路径和后端中编译计算属性提供了一个共同的 API。

参数：

• args_info (Any)
  
• out_tree (PyTreeDef)
  

__call__(*args, **kwargs)

将自身作为函数调用。

as_text()

这是可执行文件的人类可读文本表示。

旨在可视化和调试。这不是有效的也不是可靠的序列化。

如果不可用，则返回None，例如基于后端、编译器或运行时。

返回类型：

str | None

cost_analysis()

执行成本估算的摘要。

旨在可视化和调试。此输出的对象是一些简单的数据结构，可以轻松打印或序列化（例如，带有数值叶的嵌套字典、列表和元组）。然而，它的结构可以是任意的：在 JAX 和 jaxlib 的不同版本甚至调用之间可能不一致。

如果不可用，则返回None，例如基于后端、编译器或运行时。

返回类型：

Any | None

property in_tree: PyTreeDef

(位置参数，关键字参数) 的树结构。

memory_analysis()

估计内存需求的摘要。

用于可视化和调试目的。由此输出的对象是一些简单的数据结构，可以轻松打印或序列化（例如嵌套的字典、列表和具有数字叶子的元组）。然而，其结构可以是任意的：在 JAX 和 jaxlib 的不同版本之间，甚至在不同调用之间可能是不一致的。

返回 None 如果不可用，例如基于后端、编译器或运行时。

返回类型：

任意 | None

runtime_executable()

此可执行对象的任意对象表示。

用于调试目的。这不是有效也不是可靠的序列化。输出不能保证在不同调用之间的一致性。

返回 None 如果不可用，例如基于后端、编译器或运行时。

返回类型：

任意 | None