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 | None) – 可选的分片，用于指定应从哪个设备调用回调。

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

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

返回:

一个 jax.Array 对象的 pytree，其结构与以下对象的结构匹配

result_shape_dtypes.

返回类型:

result

另请参阅