[keras/src/print.py] Implement backend-specialised print function

keras/api/__init__.py

@@ -60,6 +60,7 @@
 from keras.src.ops.function import Function as Function
 from keras.src.ops.operation import Operation as Operation
 from keras.src.optimizers.optimizer import Optimizer as Optimizer
+from keras.src.print import print as print
 from keras.src.quantizers.quantizers import Quantizer as Quantizer
 from keras.src.regularizers.regularizers import Regularizer as Regularizer
 from keras.src.version import __version__ as __version__

keras/api/_tf_keras/keras/__init__.py

@@ -58,6 +58,7 @@
 from keras.src.ops.function import Function as Function
 from keras.src.ops.operation import Operation as Operation
 from keras.src.optimizers.optimizer import Optimizer as Optimizer
+from keras.src.print import print as print
 from keras.src.quantizers.quantizers import Quantizer as Quantizer
 from keras.src.regularizers.regularizers import Regularizer as Regularizer
 from keras.src.version import __version__ as __version__

keras/src/__init__.py

@@ -17,4 +17,5 @@
 from keras.src.models import Functional
 from keras.src.models import Model
 from keras.src.models import Sequential
+from keras.src.print import print
 from keras.src.version import __version__

keras/src/print.py

@@ -0,0 +1,25 @@
+from keras.src.api_export import keras_export
+from keras.src.backend import backend as keras_backend
+
+_print = print
+
+
+@keras_export("keras.print")
+def print(*args, **kwargs):
+    backend = keras_backend()
+    if backend == "jax":
+        import jax  # noqa: E402
+
+        print_fn = jax.debug.print
+    elif backend == "tensorflow":
+        import tensorflow as tf  # noqa: E402
+
+        print_fn = tf.print
+    else:
+        print_fn = _print
+    # TODO:
+    # "torch"
+    #   pytorch.org/docs/stable/generated/torch.set_printoptions.html ?
+    # "openvino"
+    # "numpy"
+    return print_fn(*args, **kwargs)