Why can't specify output stream to jax.debug.print?

[cool-RR]

I want to use jax.debug.print to debug by code, but Jax is being run by third-party code that takes over stdout. This means that the output of jax.debug.print is not shown. I'd like to direct that stream to write to a file, but I don't see that jax.debug.print has an argument for that. Is there a reason for that? Otherwise I can create an issue for this feature request.

[jakevdp]

debug.print is a very thin wrapper around debug.callback, as mentioned in the doc: https://jax.readthedocs.io/en/latest/_autosummary/jax.debug.print.html#jax.debug.print

It's deliberately kept very simple. If you want more flexibility in what is printed or how it is printed, I'd suggest using debug.callback directly.