jax.debug.print

内容

jax.debug.print#

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

打印值并在分阶段的 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

内容