JAX 调试标志#

JAX 提供了标志和上下文管理器,可以更轻松地捕获错误。

jax_debug_nans 配置选项和上下文管理器#

摘要: 启用 jax_debug_nans 标志以自动检测在 jax.jit 编译的代码中何时产生 NaN(但不在 jax.pmapjax.pjit 编译的代码中)。

jax_debug_nans 是一个 JAX 标志,启用后,当检测到 NaN 时会自动引发错误。它对 JIT 编译有特殊的处理 - 当从 JIT 函数检测到 NaN 输出时,该函数会急切地重新运行(即不进行编译),并在产生 NaN 的特定原语处抛出错误。

用法#

如果您想追踪 NaN 在您的函数或梯度中出现的位置,您可以通过以下方式开启 NaN 检查器:

  • 设置 JAX_DEBUG_NANS=True 环境变量;

  • 在您的主文件顶部附近添加 jax.config.update("jax_debug_nans", True)

  • jax.config.parse_flags_with_absl() 添加到您的主文件,然后使用命令行标志(如 --jax_debug_nans=True)设置选项;

示例#

import jax
jax.config.update("jax_debug_nans", True)

def f(x, y):
  return x / y
jax.jit(f)(0., 0.)  # ==> raises FloatingPointError exception!

jax_debug_nans 的优势和局限性#

优势#
  • 易于应用

  • 精确检测 NaN 的产生位置

  • 抛出标准的 Python 异常,并与 PDB 事后调试兼容

局限性#
  • jax.pmapjax.pjit 不兼容

  • 急切地重新运行函数可能会很慢

  • 对误报(例如,有意创建的 NaN)报错

jax_disable_jit 配置选项和上下文管理器#

摘要: 启用 jax_disable_jit 标志以禁用 JIT 编译,从而可以使用传统的 Python 调试工具,如 printpdb

jax_disable_jit 是一个 JAX 标志,启用后,会在整个 JAX 中禁用 JIT 编译(包括在控制流函数中,如 jax.lax.condjax.lax.scan)。

用法#

您可以通过以下方式禁用 JIT 编译:

  • 设置 JAX_DISABLE_JIT=True 环境变量;

  • 在您的主文件顶部附近添加 jax.config.update("jax_disable_jit", True)

  • jax.config.parse_flags_with_absl() 添加到您的主文件,然后使用命令行标志(如 --jax_disable_jit=True)设置选项;

示例#

import jax
jax.config.update("jax_disable_jit", True)

def f(x):
  y = jnp.log(x)
  if jnp.isnan(y):
    breakpoint()
  return y
jax.jit(f)(-2.)  # ==> Enters PDB breakpoint!

jax_disable_jit 的优势和局限性#

优势#
  • 易于应用

  • 启用 Python 内置的 breakpointprint 的使用

  • 抛出标准的 Python 异常,并与 PDB 事后调试兼容

局限性#
  • jax.pmapjax.pjit 不兼容

  • 在没有 JIT 编译的情况下运行函数可能会很慢