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) – 一个仅关键字参数，用于指示分阶段计算是否将强制此 jax.debug.breakpoint 与其他有序的 jax.debug.breakpoint 和 jax.debug.print 调用保持顺序。

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

返回:

如果传递了 token，则其值将原样返回。否则，返回 None。