jax.experimental.host_callback.id_print

内容

jax.experimental.host_callback.id_print#

jax.experimental.host_callback.id_print(arg, *, result=None, tap_with_device=False, device_index=0, output_stream=None, threshold=None, callback_flavor=CallbackFlavor.IO_CALLBACK, **kwargs)[source]#

类似于 id_tap()，但使用打印 tap 函数。

警告

从 2024 年 3 月 20 日起，host_callback API 已被弃用。其功能被 新的 JAX 外部回调 包含。请查看 google/jax#20385。

在每次调用打印 tap 时，如果存在 kwargs，将首先打印它（按键排序）。然后打印 arg，数组将使用 numpy.array2string 字符串化。

参见 id_tap() 文档。

附加关键字参数

  • tap_with_device 如果为 True，还将打印值来源的设备。

  • output_stream 如果给出，将用它代替内置的 print。字符串将作为 output_stream.write(s) 传递。

  • threshold 传递给 numpy.array2string。

  • callback_flavor: 如果使用 JAX_HOST_CALLBACK_LEGACY=False 运行，则指定

    要使用的回调类型。参见 google/jax#20385.

有关更多详细信息，请参见 jax.experimental.host_callback 模块文档。

目录