jax.debug.callback#

jax.debug.callback(callback, *args, ordered=False, partitioned=False, **kwargs)[source]#

调用一个可分阶段的 Python 回调。

欲了解更多信息，请参阅外部回调。

jax.debug.callback 允许您传入一个可在 JAX 阶段化程序中调用的 Python 函数。jax.debug.callback 遵循现有的 JAX 转换的纯操作语义，因此不感知副作用。这意味着在存在高阶原语和转换的情况下，效果可能会被丢弃、复制或潜在地重新排序。

我们希望这种行为，因为我们希望 jax.debug.callback 是“无害的”，也就是说，我们希望这些原语尽可能少地改变 JAX 计算，同时尽可能多地揭示有关它们的信息，例如计算的哪些部分被复制或丢弃。

参数:

  • callback (Callable[..., None]) – 一个返回 None 的 Python 可调用对象。

  • *args (Any) – 回调函数的位置参数。

  • ordered (bool) – 一个仅限关键字的参数，用于指示阶段化计算是否会强制此回调与其他有序回调之间的顺序。

  • partitioned (bool) – 如果为 True，则仅打印本地分片；此选项可避免操作数进行全收集。如果为 False，则与逻辑操作数一起打印；此选项首先需要操作数进行全收集。

  • **kwargs (Any) – 回调函数的关键字参数。

返回:

无

返回类型:

无

另请参阅