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

keras/api/_tf_keras/keras/ops/__init__.py

@@ -18,6 +18,7 @@
 from keras.src.ops.core import fori_loop as fori_loop
 from keras.src.ops.core import is_tensor as is_tensor
 from keras.src.ops.core import map as map
+from keras.src.ops.core import print as print
 from keras.src.ops.core import saturate_cast as saturate_cast
 from keras.src.ops.core import scan as scan
 from keras.src.ops.core import scatter as scatter

keras/api/ops/__init__.py

@@ -18,6 +18,7 @@
 from keras.src.ops.core import fori_loop as fori_loop
 from keras.src.ops.core import is_tensor as is_tensor
 from keras.src.ops.core import map as map
+from keras.src.ops.core import print as print
 from keras.src.ops.core import saturate_cast as saturate_cast
 from keras.src.ops.core import scan as scan
 from keras.src.ops.core import scatter as scatter

keras/src/backend/jax/core.py

@@ -428,3 +428,13 @@ def device_scope(device_name):
     else:
         jax_device = device_name
     return jax.default_device(jax_device)
+
+
+def print(*args, **kwargs):
+    """Prints values and works in staged out JAX functions.
+
+    This function does *not* work with f-strings because formatting is delayed.
+    So instead of ``jax.debug.print(f"hello {bar}")``, write
+    ``jax.debug.print("hello {bar}", bar=bar)``.
+    """
+    return jax.debug.print(*args, **kwargs)

keras/src/backend/tensorflow/core.py

@@ -696,3 +696,14 @@ def __exit__(self, *args, **kwargs):
 
 def device_scope(device_name):
     return tf.device(device_name)
+
+
+def print(*args, **kwargs):
+    """Print the specified inputs.
+
+    A TensorFlow operator that prints the specified inputs to a desired
+    output stream or logging level. The inputs may be dense or sparse Tensors,
+    primitive python objects, data structures that contain tensors, and
+    printable Python objects. Printed tensors will recursively show the first
+    and last elements of each dimension to summarize."""
+    return tf.print(*args, **kwargs)

keras/src/ops/core.py

@@ -1183,3 +1183,13 @@ def grad(*args, upstream):
     ```
     """
     return backend.core.custom_gradient(f)
+
+
+_print = print
+
+
+@keras_export("keras.ops.print")
+def print(*args, **kwargs):
+    return (backend.core.print if hasattr(backend.core, "print") else _print)(
+        *args, **kwargs
+    )