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) – 可选的分片（sharding），指定应从中调用回调函数的设备。

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

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

返回:

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

result_shape_dtypes.

返回类型:

result

另请参阅