keras.ops.print: tf.print; jax.debug.print; … #21137
Description
As per our discussion earlier today @fchollet
https://www.tensorflow.org/api_docs/python/tf/print
https://docs.jax.dev/en/latest/_autosummary/jax.debug.print.html
Where do you think this function should be? keras.utils.summary_utils? keras.ops.print? - Somewhere else?
Roughly:
from keras.api import backend
import jax
import tensorflow as tf
_print = print
def print(*args, **kwargs):
backend = backend.backend()
print_fn = {"jax": jax.debug.print,
"tensorflow": tf.print}.get(backend, _print)
# "torch" https://pytorch.org/docs/stable/generated/torch.set_printoptions.html ?
# "openvino"
# "numpy"
return print_fn(*args, **kwargs)