Calls an impure Python callback.
For more explanation, see External Callbacks.
callback (Callable[..., Any]) – function to execute on the host. It is assumed 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. jax.ShapeDtypeStruct is often used to define leaf values.
*args (Any) – arguments to be passed to the callback function
sharding (SingleDeviceSharding | None | 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
jax.Array objects whose structure matches that of
result_shape_dtypes.
result
RetroSearch is an open source project built by @garambo | Open a GitHub Issue
Search and Browse the WWW like it's 1997 | Search results from DuckDuckGo
HTML:
3.2
| Encoding:
UTF-8
| Version:
0.7.3