jax.debug.breakpoint#

jax.debug.breakpoint(*, backend=None, filter_frames=True, num_frames=None, ordered=False, token=None, **kwargs)[source]#

在程序中的某个点进入断点。

参数:

  • backend (str | None | None) – 要使用的调试器后端。默认情况下，选择优先级最高的调试器，在没有其他已注册调试器的情况下，回退到 CLI 调试器。

  • filter_frames (bool) – 是否从回溯中过滤掉 JAX 内部堆栈帧。由于某些库（如 Flax）也使用 JAX 的堆栈帧过滤系统，此选项也可能影响是否过滤库中的堆栈帧。

  • num_frames (int | None | None) – 在交互式调试器中可用于检查的当前堆栈帧之上的帧数。

  • ordered (bool) – 仅关键字参数，用于指示 staged out 计算是否将强制执行此 jax.debug.breakpoint 相对于其他 ordered jax.debug.breakpoint 和 jax.debug.print 调用的排序。

  • token – 仅关键字参数；ordered 的替代方案。如果使用，则应传递 JAX 数组（或 JAX 数组的 pytree），并且断点将在计算出其值后运行。这会原封不动地返回，并且应该传递回计算。如果在后续计算中未使用返回值，则将修剪整个计算，并且不会运行此断点。

返回:

如果传递了 token，则其值将原封不动地返回。否则，返回 None。