jax.experimental.io_callback#

jax.experimental.io_callback(callback, result_shape_dtypes, *args, sharding=None, ordered=False, **kwargs)[来源]#

调用一个不纯的 Python 回调。

有关更多解释，请参阅 外部回调。

参数:

  • callback (Callable[..., Any]) – 要在主机上执行的函数。假定它是一个不纯函数。如果 callback 是纯函数，则使用 jax.pure_callback() 可能导致更高效的执行。

  • result_shape_dtypes (Any) – pytree，其叶子具有 shape 和 dtype 属性，其结构与运行时回调函数的预期输出匹配。 jax.ShapeDtypeStruct 通常用于定义叶子值。

  • *args (Any) – 要传递给回调函数的参数

  • sharding (SingleDeviceSharding | None) – 可选的分片，用于指定应从中调用回调的设备。

  • ordered (bool) – 布尔值，指定对回调的顺序调用是否必须是有序的。

  • **kwargs (Any) – 要传递给回调函数的关键字参数

返回:

一个 jax.Array 对象的 pytree，其结构与

result_shape_dtypes.

返回类型:

result

另请参阅