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)[源代码]#

类似于带有打印功能的 id_tap()。

警告

自2024年3月20日起，host_callback API 已被弃用。该功能已被 新的 JAX 外部回调 所取代。详情请参见 google/jax#20385。

在每次调用打印水龙头时，如果存在 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 模块文档。

目录