JAX 中文文档（十三）（3）

JAX 中文文档（十三）（2）https://developer.aliyun.com/article/1559742

GPU 内存分配

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

当第一个 JAX 操作运行时，JAX 将预先分配总 GPU 内存的 75%。 预先分配可以最小化分配开销和内存碎片化，但有时会导致内存不足（OOM）错误。如果您的 JAX 进程因内存不足而失败，可以使用以下环境变量来覆盖默认行为：

XLA_PYTHON_CLIENT_PREALLOCATE=false

这将禁用预分配行为。JAX 将根据需要分配 GPU 内存，可能会减少总体内存使用。但是，这种行为更容易导致 GPU 内存碎片化，这意味着使用大部分可用 GPU 内存的 JAX 程序可能会在禁用预分配时发生 OOM。

XLA_PYTHON_CLIENT_MEM_FRACTION=.XX

如果启用了预分配，这将使 JAX 预分配总 GPU 内存的 XX% ，而不是默认的 75%。减少预分配量可以修复 JAX 程序启动时的内存不足问题。

XLA_PYTHON_CLIENT_ALLOCATOR=platform

这使得 JAX 根据需求精确分配内存，并释放不再需要的内存（请注意，这是唯一会释放 GPU 内存而不是重用它的配置）。这样做非常慢，因此不建议用于一般用途，但可能对于以最小可能的 GPU 内存占用运行或调试 OOM 失败非常有用。

OOM 失败的常见原因

同时运行多个 JAX 进程。

要么使用 XLA_PYTHON_CLIENT_MEM_FRACTION 为每个进程分配适当的内存量，要么设置 XLA_PYTHON_CLIENT_PREALLOCATE=false。

同时运行 JAX 和 GPU TensorFlow。

TensorFlow 默认也会预分配，因此这与同时运行多个 JAX 进程类似。

一个解决方案是仅使用 CPU TensorFlow（例如，如果您仅使用 TF 进行数据加载）。您可以使用命令 tf.config.experimental.set_visible_devices([], "GPU") 阻止 TensorFlow 使用 GPU。

或者，使用 XLA_PYTHON_CLIENT_MEM_FRACTION 或 XLA_PYTHON_CLIENT_PREALLOCATE。还有类似的选项可以配置 TensorFlow 的 GPU 内存分配（gpu_memory_fraction 和 allow_growth 在 TF1 中应该设置在传递给 tf.Session 的 tf.ConfigProto 中。参见 使用 GPU：限制 GPU 内存增长 用于 TF2）。

在显示 GPU 上运行 JAX。

使用 XLA_PYTHON_CLIENT_MEM_FRACTION 或 XLA_PYTHON_CLIENT_PREALLOCATE。

提升秩警告

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

NumPy 广播规则 允许自动将参数从一个秩（数组轴的数量）提升到另一个秩。当意图明确时，此行为很方便，但也可能导致意外的错误，其中静默的秩提升掩盖了潜在的形状错误。

下面是提升秩的示例：

>>> import numpy as np
>>> x = np.arange(12).reshape(4, 3)
>>> y = np.array([0, 1, 0])
>>> x + y
array([[ 0,  2,  2],
 [ 3,  5,  5],
 [ 6,  8,  8],
 [ 9, 11, 11]]) 

为了避免潜在的意外，jax.numpy 可配置，以便需要提升秩的表达式会导致警告、错误或像常规 NumPy 一样允许。配置选项名为 jax_numpy_rank_promotion，可以取字符串值 allow、warn 和 raise。默认设置为 allow，允许提升秩而不警告或错误。设置为 raise 则在提升秩时引发错误，而 warn 在首次提升秩时引发警告。

可以使用 jax.numpy_rank_promotion() 上下文管理器在本地启用或禁用提升秩：

with jax.numpy_rank_promotion("warn"):
  z = x + y 

这个配置也可以在多种全局方式下设置。其中一种是在代码中使用 jax.config：

import jax
jax.config.update("jax_numpy_rank_promotion", "warn") 

也可以使用环境变量 JAX_NUMPY_RANK_PROMOTION 来设置选项，例如 JAX_NUMPY_RANK_PROMOTION='warn'。最后，在使用 absl-py 时，可以使用命令行标志设置选项。

公共 API：jax 包

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

子包

• 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 模块
  

配置

即时编译 (jit)

| block_until_ready(x) | 尝试调用 pytree 叶子上的 block_until_ready 方法。 | ## 自动微分

jax.Array (jax.Array)

向量化 (vmap)

并行化 (pmap)

Callbacks

Miscellaneous

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