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

SamuelMarks · 2025-06-02T20:10:07Z

… run --all-files`

codecov-commenter · 2025-06-02T20:16:53Z

Codecov Report

Attention: Patch coverage is 61.11111% with 7 lines in your changes missing coverage. Please review.

Project coverage is 82.67%. Comparing base (ad7bbb8) to head (6c7db62).
Report is 1 commits behind head on master.

Files with missing lines	Patch %	Lines
keras/src/backend/numpy/core.py	50.00%	2 Missing ⚠️
keras/src/backend/torch/core.py	50.00%	2 Missing ⚠️
keras/api/_tf_keras/keras/ops/__init__.py	0.00%	1 Missing ⚠️
keras/src/backend/openvino/core.py	66.66%	1 Missing ⚠️
keras/src/ops/core.py	66.66%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #21344      +/-   ##
==========================================
- Coverage   82.68%   82.67%   -0.01%     
==========================================
  Files         565      565              
  Lines       54829    54847      +18     
  Branches     8513     8513              
==========================================
+ Hits        45335    45346      +11     
- Misses       7409     7416       +7     
  Partials     2085     2085

Flag	Coverage Δ
keras	`82.48% <61.11%> (-0.01%)`	⬇️
keras-jax	`63.62% <38.88%> (-0.01%)`	⬇️
keras-numpy	`58.77% <50.00%> (-0.01%)`	⬇️
keras-openvino	`33.12% <38.88%> (+<0.01%)`	⬆️
keras-tensorflow	`64.01% <38.88%> (-0.01%)`	⬇️
keras-torch	`63.65% <38.88%> (-0.01%)`	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

fchollet

Thanks for the PR!

If the op is not implemented for a given backend, create the op anyway and raise NotImplementedError.

Also please add a simple unit test.

fchollet · 2025-06-02T21:18:46Z

keras/src/print.py

+_print = print
+
+
+@keras_export("keras.print")


Instead, please add individual print ops to e.g. keras/src/backend/tensorflow/core.py, keras/src/backend/jax/core.py, etc. Then export the op is keras/src/ops/core.py.

Ok I'm not sure if you want from jax.debug import print and from tensorflow import print or if you want them with first-class docstrings. I've added the latter; and updated this PR.

I believe the request was to add a print implementation also for the torch, numpy and openvino backends (even if they're the same).

Then remove the fallback mechanism and _print.

…on implementation across backend-specific internal modules

hertschuh · 2025-06-03T21:25:04Z

keras/src/backend/jax/core.py

+
+
+def print(*args, **kwargs):
+    """Prints values and works in staged out JAX functions.


The backend specific docstring is not visible anywhere (neither on keras.io, nor in an IDE). So move to the docstring of keras/src/ops/core.py explaining that this is JAX specific.

hertschuh · 2025-06-03T21:26:29Z

keras/src/backend/tensorflow/core.py

+
+
+def print(*args, **kwargs):
+    """Print the specified inputs.


The backend specific docstring is not visible anywhere (neither on keras.io, nor in an IDE). So if you think it is useful, move to the docstring of keras/src/ops/core.py explaining that this is TensorFlow specific.

hertschuh · 2025-06-03T21:26:54Z

keras/src/backend/jax/core.py

+    """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


Single back ticks ` here and on the next line

hertschuh · 2025-06-03T21:29:22Z

keras/src/print.py

+_print = print
+
+
+@keras_export("keras.print")


I believe the request was to add a print implementation also for the torch, numpy and openvino backends (even if they're the same).

Then remove the fallback mechanism and _print.

hertschuh

I'm having second thoughts about this. In its current form, keras.ops.print is not cross-backend in the way we want it. Take this code:

x = keras.ops.array(1.0)
keras.ops.print("x is", x)

This will work on all backends except JAX, it will fail on JAX.

However:

x = keras.ops.array(1.0)
keras.ops.print("x is {x}", x=x)

Will work on JAX, but fail on all other backends.

In fact there is no way to print text and a tensor that will work on JAX and all other backends. This should be a minimal requirement for this feature: have a truly cross backend way to print text and tensors.

This means that the JAX implementation has to be changed. It is probably doable fairly easily using jax.debug.callback instead of jax.debug.print.

Then there is the question of keyword arguments. Python's print has:

sep
end
file
flush

Tensorflow supports:

sep
end
output_stream

Instead of **kwargs, I think we should explicitly spell out sep=' ', end='\n', file=None, flush=False.

On Tensorflow, flush will be ignored and file will be remapped to output_stream.
On Torch, Numpy, OpenVino, things will just work with Python's print
On JAX, the new implementation needs to support them

hertschuh · 2025-06-04T17:58:01Z

keras/src/backend/numpy/core.py

+
+
+def print(*args, print_options=None, **kwargs):
+    np.set_printoptions(


The issue is that np.set_printoptions is global, not just for the scope of this call to print.

Let's say somebody does:

np.set_printoptions(threshold=100) ... keras.ops.print(x)

They would never understand why their setting of 100 is not working.

So let's remove this.

hertschuh · 2025-06-04T17:58:45Z

keras/src/backend/numpy/core.py

@@ -452,3 +454,10 @@ def remat(f):
        "utilize this feature."
    )
    return f
+
+
+def print(*args, print_options=None, **kwargs):


The signature needs to be consistent between all the backends, so remove print_options.

hertschuh · 2025-06-04T18:00:19Z

keras/src/backend/numpy/core.py

@@ -18,6 +18,8 @@
 SUPPORTS_RAGGED_TENSORS = False
 IS_THREAD_SAFE = True

+_print = print


nitpick: instead of doing this, let's use __builtins__:

def print(*args, **kwargs): return __builtins__.print(*args, **kwargs)

Oh that's a good one, I didn't know the __builtins__

hertschuh · 2025-06-04T18:02:29Z

keras/src/backend/jax/core.py

@@ -1,3 +1,4 @@
+from jax.debug import print  # noqa


Let's be explicit instead, even if it seems redundant:

def print(*args, **kwargs): jax.debug.print(*args, **kwargs)

That being said, per my overall comment, the JAX implementation needs to be completely different and not use jax.debug.print.

Hold up a moment I'll send JAX a PR; at latest this weekend

Wait, I don't think we should change JAX. I believe we can simply do:

def print(*args, sep=' ', end='\n', file=None, flush=False): kwargs = {"sep": sep, "end": end, "file": file, "flush": flush} jax.debug.callback( lambda *args, **kwargs: print(*args, **kwargs), *args, **kwargs)

But what about that whole interpolation discrepancy you referenced above?

I guess, for consistency, we're using the TF style, which is also the normal Python print style without f-string or format.

So we should document that: you cannot use f-string or format, you have to pass tensors as argument, and give an example: keras.ops.print("x:", x, "y:", y).

Hold on let me send JAX a PR and see what they say

hertschuh · 2025-06-04T18:02:49Z

keras/src/backend/openvino/core.py

@@ -22,6 +22,8 @@
 SUPPORTS_RAGGED_TENSORS = False
 IS_THREAD_SAFE = True

+_print = print


nitpick: instead of doing this, let's use __builtins__:

def print(*args, **kwargs): return __builtins__.print(*args, **kwargs)

hertschuh · 2025-06-04T18:03:04Z

keras/src/backend/torch/core.py

@@ -23,6 +23,8 @@
 SUPPORTS_RAGGED_TENSORS = False
 IS_THREAD_SAFE = True

+_print = print


nitpick: instead of doing this, let's use __builtins__:

def print(*args, **kwargs): return __builtins__.print(*args, **kwargs)

hertschuh · 2025-06-04T18:04:05Z

keras/src/backend/torch/core.py

+
+
+def print(*args, print_options=None, **kwargs):
+    torch.set_printoptions(


The issue is that torch.set_printoptions is global, not just for the scope of this call to print.

Let's say somebody does:

torch.set_printoptions(threshold=100) ... keras.ops.print(x)

They would never understand why their setting of 100 is not working.

So let's remove this.

hertschuh · 2025-06-04T18:05:37Z

keras/src/ops/core.py

+
+@keras_export("keras.ops.print")
+def print(*args, **kwargs):
+    """Backend-specialised print function, oft handles tensors and


Docstring needs to have the following format:

"""One-line short description. Optionally more paragraphs of details. """

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

d31c128

google-ml-butler bot added the size:S label Jun 2, 2025

google-ml-butler bot assigned gbaned Jun 2, 2025

[keras/src/print.py] Use keras.src.backend import ; [*] `pre-commit…

4571fee

… run --all-files`

fchollet reviewed Jun 2, 2025

View reviewed changes

[keras/src/ops/core.py] Refactor to put print in ops and distributi…

726837c

…on implementation across backend-specific internal modules

hertschuh requested changes Jun 3, 2025

View reviewed changes

[keras/src/backend/*/core.py] Add print impl for each backend

6c7db62

hertschuh reviewed Jun 4, 2025

View reviewed changes



		def print(args, *kwargs):
		"""Prints values and works in staged out JAX functions.



		def print(args, print_options=None, *kwargs):
		np.set_printoptions(



		def print(args, print_options=None, *kwargs):
		torch.set_printoptions(

		_print = print


		@keras_export("keras.print")

		_print = print


		@keras_export("keras.print")

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

Are you sure you want to change the base?

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

Uh oh!

Conversation

SamuelMarks commented Jun 2, 2025

Uh oh!

codecov-commenter commented Jun 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

fchollet left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

hertschuh left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

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

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

codecov-commenter commented Jun 2, 2025 •

edited

Loading