jax.experimental.pallas.debug_print

jax.experimental.pallas.debug_print(fmt, *args)

在 Pallas 内核内部打印值。

参数:

• fmt (str) –

  一个包含在输出中的格式化字符串。格式化字符串的限制取决于后端。

  • 在 GPU 上，当使用 Triton 时，fmt 不得包含任何占位符 ({...})，因为它总是在任何值之前打印。

  • 在 GPU 上，当使用实验性的 Mosaic GPU 后端时，fmt 必须为每个要打印的值包含一个占位符。不支持格式规范和转换。所有值都必须是标量。

  • 在 TPU 中，如果 fmt 包含占位符，则所有值都必须是 32 位整数。如果没有占位符，则在格式化字符串之后打印这些值。所有值都必须是标量。

• *args (jax.typing.ArrayLike) – 要打印的值。