jax.experimental.io_callback#

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

调用一个不纯的 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

另请参阅