jax.experimental.host_callback.id_print¶

jax.experimental.host_callback.id_print(arg, *, result=None, tap_with_device=False, output_stream=None, threshold=None, **kwargs)[source]¶

Like id_tap() with a printing tap function.

Experimental: please give feedback, and expect changes!

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.