调试简介#

本节将向您介绍一组内置的 JAX 调试方法 —— jax.debug.print()、jax.debug.breakpoint() 和 jax.debug.callback() —— 您可以将它们与各种 JAX 转换一起使用。

让我们从 jax.debug.print() 开始。

jax.debug.print 用于简单检查#

这是一个经验法则

回想一下即时编译，当使用 jax.jit() 转换函数时，Python 代码将使用抽象追踪器代替您的数组执行。因此，Python print() 函数只会打印此追踪器值。

import jax
import jax.numpy as jnp

@jax.jit
def f(x):
  print("print(x) ->", x)
  y = jnp.sin(x)
  print("print(y) ->", y)
  return y

result = f(2.)

print(x) -> Traced<ShapedArray(float32[], weak_type=True)>with<DynamicJaxprTrace>
print(y) -> Traced<ShapedArray(float32[], weak_type=True)>with<DynamicJaxprTrace>

Python 的 print 在运行时值存在之前，在追踪时执行。如果您想打印实际的运行时值，可以使用 jax.debug.print()。

@jax.jit
def f(x):
  jax.debug.print("jax.debug.print(x) -> {x}", x=x)
  y = jnp.sin(x)
  jax.debug.print("jax.debug.print(y) -> {y}", y=y)
  return y

result = f(2.)

jax.debug.print(x) -> 2.0
jax.debug.print(y) -> 0.9092974066734314

类似地，在 jax.vmap() 中，使用 Python 的 print 只会打印追踪器；要打印正在映射的值，请使用 jax.debug.print()。

def f(x):
  jax.debug.print("jax.debug.print(x) -> {}", x)
  y = jnp.sin(x)
  jax.debug.print("jax.debug.print(y) -> {}", y)
  return y

xs = jnp.arange(3.)

result = jax.vmap(f)(xs)

jax.debug.print(x) -> 0.0
jax.debug.print(x) -> 1.0
jax.debug.print(x) -> 2.0
jax.debug.print(y) -> 0.0
jax.debug.print(y) -> 0.8414709568023682
jax.debug.print(y) -> 0.9092974066734314

这是使用 jax.lax.map() 的结果，它是一个顺序映射而不是矢量化。

result = jax.lax.map(f, xs)

jax.debug.print(x) -> 0.0
jax.debug.print(y) -> 0.0
jax.debug.print(x) -> 1.0
jax.debug.print(y) -> 0.8414709568023682
jax.debug.print(x) -> 2.0
jax.debug.print(y) -> 0.9092974066734314

请注意，顺序是不同的，因为 jax.vmap() 和 jax.lax.map() 以不同的方式计算相同的结果。在调试时，评估顺序的详细信息正是您可能需要检查的内容。

下面是一个使用 jax.grad() 的示例，其中 jax.debug.print() 仅打印前向传递。在这种情况下，行为类似于 Python 的 print()，但是如果在调用期间应用 jax.jit()，则行为是一致的。

def f(x):
  jax.debug.print("jax.debug.print(x) -> {}", x)
  return x ** 2

result = jax.grad(f)(1.)

jax.debug.print(x) -> 1.0

有时，当参数彼此不依赖时，使用 JAX 转换进行分阶段处理时，对 jax.debug.print() 的调用可能会以不同的顺序打印它们。如果您需要原始顺序，例如先 x: ...，然后是 y: ...，请添加 ordered=True 参数。

例如

@jax.jit
def f(x, y):
  jax.debug.print("jax.debug.print(x) -> {}", x, ordered=True)
  jax.debug.print("jax.debug.print(y) -> {}", y, ordered=True)
  return x + y

f(1, 2)

jax.debug.print(x) -> 1
jax.debug.print(y) -> 2

Array(3, dtype=int32, weak_type=True)

要了解有关 jax.debug.print() 及其 Sharp Bits 的更多信息，请参阅高级调试。

jax.debug.breakpoint 用于类似 pdb 的调试#

总结：使用 jax.debug.breakpoint() 暂停 JAX 程序的执行以检查值。

要在调试期间在某些点暂停已编译的 JAX 程序，可以使用 jax.debug.breakpoint()。提示类似于 Python pdb，它允许您检查调用堆栈中的值。实际上，jax.debug.breakpoint() 是 jax.debug.callback() 的一个应用，它捕获有关调用堆栈的信息。

要在 breakpoint 调试会话期间打印所有可用命令，请使用 help 命令。（完整的调试器命令、Sharp Bits、其优点和局限性在高级调试中介绍。）

这是一个调试器会话可能如下所示的示例

@jax.jit
def f(x):
  y, z = jnp.sin(x), jnp.cos(x)
  jax.debug.breakpoint()
  return y * z
f(2.) # ==> Pauses during execution

JAX debugger

对于依赖于值的断点，您可以使用运行时条件，如 jax.lax.cond()

def breakpoint_if_nonfinite(x):
  is_finite = jnp.isfinite(x).all()
  def true_fn(x):
    pass
  def false_fn(x):
    jax.debug.breakpoint()
  jax.lax.cond(is_finite, true_fn, false_fn, x)

@jax.jit
def f(x, y):
  z = x / y
  breakpoint_if_nonfinite(z)
  return z

f(2., 1.) # ==> No breakpoint

Array(2., dtype=float32, weak_type=True)

f(2., 0.) # ==> Pauses during execution

jax.debug.callback 用于调试期间的更多控制#

jax.debug.print() 和 jax.debug.breakpoint() 都是使用更灵活的 jax.debug.callback() 实现的，它可以通过 Python 回调更好地控制主机端逻辑的执行。它与 jax.jit()、jax.vmap()、jax.grad() 和其他转换兼容（有关更多信息，请参阅回调的风格 表格中的 外部回调）。

例如

import logging

def log_value(x):
  logging.warning(f'Logged value: {x}')

@jax.jit
def f(x):
  jax.debug.callback(log_value, x)
  return x

f(1.0);

WARNING:root:Logged value: 1.0

此回调与其他转换兼容，包括 jax.vmap() 和 jax.grad()。

x = jnp.arange(5.0)
jax.vmap(f)(x);

WARNING:root:Logged value: 0.0
WARNING:root:Logged value: 1.0
WARNING:root:Logged value: 2.0
WARNING:root:Logged value: 3.0
WARNING:root:Logged value: 4.0

jax.grad(f)(1.0);

WARNING:root:Logged value: 1.0

这使得 jax.debug.callback() 可用于通用调试。

您可以在外部回调中了解有关 jax.debug.callback() 和其他 JAX 回调的更多信息。

下一步#

查看高级调试以了解有关在 JAX 中进行调试的更多信息。