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 位元整數。如果沒有預留位置，則值會在格式字串之後列印。

  • 在 TPU 上，如果輸入是單一向量，則向量會在格式字串之後列印。格式字串必須以單一預留位置 {} 結尾。

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