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

Calls an impure Python callback.

For more explanation, see External Callbacks.

  • callback (Callable[..., Any]) – function to execute on the host. It is assumet to be an impure function. If callback is pure, using jax.pure_callback() instead may lead to more efficient execution.

  • result_shape_dtypes (Any) – pytree whose leaves have shape and dtype attributes, whose structure matches the expected output of the callback function at runtime.

  • *args (Any) – arguments to be passed to the callback function

  • sharding (SingleDeviceSharding | None) – optional sharding that specifies the device from which the callback should be invoked.

  • ordered (bool) – boolean specifying whether sequential calls to callback must be ordered.

  • **kwargs (Any) – keyword arguments to be passed to the callback function


a pytree of jax.Array objects whose structure matches that of


Return type


See also