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.


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 with numpy.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-in print. The string will be passed as output_stream.write(s).

  • threshold is passed to numpy.array2string.

  • callback_flavor: if running with JAX_HOST_CALLBACK_LEGACY=False specifies

    the flavor of callback to use. See google/jax#20385.

For more details see the jax.experimental.host_callback module documentation.