jax.debug.print

jax.debug.print(fmt, *args, ordered=False, **kwargs)

打印值，并可在分阶段输出的 JAX 函数中使用。

此函数不适用于 f 字符串，因为格式化是延迟的。因此，请不要使用 jax.debug.print(f"hello {bar}")，而是使用 jax.debug.print("hello {bar}", bar=bar)。

此函数是 jax.debug.callback() 的一个简单的便捷包装器。其实现本质上是

def debug_print(fmt: str, *args, **kwargs):
  jax.debug.callback(
      lambda *args, **kwargs: print(fmt.format(*args, **kwargs)),
      *args, **kwargs)

直接调用 jax.debug.callback() 而不是使用此便捷包装器可能很有用。例如，要在日志中进行调试打印，您可以使用 jax.debug.callback() 与 logging.log 一起使用。

参数:

• fmt ( str ) – 一个格式字符串，例如 "hello {x}"，它将用于格式化输入参数，类似于 str.format。请参阅关于 字符串格式化 和 格式字符串语法 的 Python 文档。

• *args – 要格式化的位置参数列表，如同传递给 fmt.format 一样。

• ordered ( bool ) – 一个仅关键字参数，用于指示分阶段输出的计算是否将强制此 jax.debug.print 相对于其他有序的 jax.debug.print 调用进行排序。

• **kwargs – 要格式化的附加关键字参数，如同传递给 fmt.format 一样。

返回类型:

None