jax.debug.breakpoint#

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

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

参数:

  • 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。