jax.experimental.host_callback.id_print

Contents

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, **kwargs)#

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

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.

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