调试运行时值#
你是否遇到了梯度爆炸?NaN 是否让你感到困扰?是否想探究计算过程中的中间值?快来了解一下以下 JAX 调试工具!本页提供了摘要,你可以点击底部的“阅读更多”链接以了解更多信息。
目录
使用 jax.debug
进行交互式检查#
完整指南请点击这里
摘要: 使用 jax.debug.print()
在 jax.jit
-、jax.pmap
- 和 pjit
- 修饰的函数中将值打印到 stdout,并使用 jax.debug.breakpoint()
暂停已编译函数的执行,以便检查调用堆栈中的值
import jax
import jax.numpy as jnp
@jax.jit
def f(x):
jax.debug.print("🤯 {x} 🤯", x=x)
y = jnp.sin(x)
jax.debug.breakpoint()
jax.debug.print("🤯 {y} 🤯", y=y)
return y
f(2.)
# Prints:
# 🤯 2.0 🤯
# Enters breakpoint to inspect values!
# 🤯 0.9092974662780762 🤯
阅读更多.
使用 jax.experimental.checkify
进行函数式错误检查#
完整指南请点击这里
摘要: Checkify 允许你向 JAX 代码添加可 jit
的运行时错误检查(例如,越界索引)。将 checkify.checkify
转换与类似断言的 checkify.check
函数一起使用,以向 JAX 代码添加运行时检查
from jax.experimental import checkify
import jax
import jax.numpy as jnp
def f(x, i):
checkify.check(i >= 0, "index needs to be non-negative!")
y = x[i]
z = jnp.sin(y)
return z
jittable_f = checkify.checkify(f)
err, z = jax.jit(jittable_f)(jnp.ones((5,)), -1)
print(err.get())
# >> index needs to be non-negative! (check failed at <...>:6 (f))
你还可以使用 checkify 自动添加常用检查
errors = checkify.user_checks | checkify.index_checks | checkify.float_checks
checked_f = checkify.checkify(f, errors=errors)
err, z = checked_f(jnp.ones((5,)), 100)
err.throw()
# ValueError: out-of-bounds indexing at <..>:7 (f)
err, z = checked_f(jnp.ones((5,)), -1)
err.throw()
# ValueError: index needs to be non-negative! (check failed at <…>:6 (f))
err, z = checked_f(jnp.array([jnp.inf, 1]), 0)
err.throw()
# ValueError: nan generated by primitive sin at <...>:8 (f)
阅读更多.
使用 JAX 的调试标志抛出 Python 错误#
完整指南请点击这里
摘要: 启用 jax_debug_nans
标志以自动检测在 jax.jit
编译的代码中何时产生 NaN(但不在 jax.pmap
或 jax.pjit
编译的代码中),并启用 jax_disable_jit
标志以禁用 JIT 编译,从而可以使用传统的 Python 调试工具,如 print
和 pdb
。
import jax
jax.config.update("jax_debug_nans", True)
def f(x, y):
return x / y
jax.jit(f)(0., 0.) # ==> raises FloatingPointError exception!
阅读更多.