jax.experimental.host_callback.id_print#
- jax.experimental.host_callback.id_print(arg, *, result=None, tap_with_device=False, device_index=0, output_stream=None, threshold=None, callback_flavor=CallbackFlavor.IO_CALLBACK, **kwargs)[source]#
Like
id_tap()
with a printing tap function.Warning
The host_callback APIs are deprecated as of March 20, 2024. The functionality is subsumed by the new JAX external callbacks See google/jax#20385.
On each invocation of the printing tap, the
kwargs
if present will be printed first (sorted by keys). Then arg will be printed, with the arrays stringified withnumpy.array2string
.See the
id_tap()
documentation.Additional keyword arguments:
tap_with_device
if True, will print also the device from which the value originates.output_stream
if given then it will be used instead of the built-inprint
. The string will be passed asoutput_stream.write(s)
.threshold
is passed tonumpy.array2string
.callback_flavor
: if running with JAX_HOST_CALLBACK_LEGACY=False specifiesthe flavor of callback to use. See google/jax#20385.
For more details see the
jax.experimental.host_callback
module documentation.