Repository: google/jax (python)

Issue:
jax.debug.print() fails when passing value formatting (which docs say is supported)
### Description

`jax.debug.print()` supports printing variable values, e.g.: `a = 5; jax.debug.print("{}", a)` outputs `5`. And this works inside of a function that has been JAX compiled as well.

And when called **without** any compiling, it supports formatting values in the standard way, e.g. `a = 5; jax.debug.print("{:.2f}", a)` outputs `5.00`.

But if you attempt to include any value formatting **in a compiled function**, then running the function **produces an error:**

```
def kernel(a):
  jax.debug.print("{:.2f}", a)
compiled = jax.jit(kernel)
compiled(5)
```

```
TypeError: unsupported format string passed to DynamicJaxprTracer.__format__
```

Yet the [documentation for jax.debug.print](https://jax.readthedocs.io/en/latest/_autosummary/jax.debug.print.html#jax.debug.print) clearly states:

> fmt (str) – A format string, e.g. "hello {x}", that will be used to format input arguments, like str.format. See the Python docs on string formatting and format string syntax.

I've encountered this bug because I'm using `jax.debug.print()` to monitor several values at a time during execution to debug it, and I'd like the rows of values to be of a similar length to make the columns more readable. Dealing with floats that might print `10.119900703430176` or might print `2` makes for a debugging experience that is harder to read.

Here is a [Colab notebook](https://colab.research.google.com/drive/1xWYGHzZHXpV10LClTyB_zYJGIVEHUBL2?usp=sharing) with example code that you can see the results and run yourself.

### System info (python version, jaxlib version, accelerator, etc.)

Running on current Google Colab.

```
jax:    0.4.26
jaxlib: 0.4.26
numpy:  1.26.4
python: 3.10.12 (main, Jul 29 2024, 16:56:48) [GCC 11.4.0]
jax.devices (1 total, 1 local): [CpuDevice(id=0)]
process_count: 1
platform: uname_result(system='Linux', node='afb8abe17318', release='6.1.85+', version='#1 SMP PREEMPT_DYNAMIC Thu Jun 27 21:05:47 UTC 2024', machine='x86_64')
```

Pull request:
Fixed validation in `jax.debug.format`
This commit ensures that no formatting is done during validation, because the arguments could be abstract values.

Closes #23475.

Patch:
diff --git a/jax/_src/debugging.py b/jax/_src/debugging.py
index 62bfdf031..3e7082ab1 100644
--- a/jax/_src/debugging.py
+++ b/jax/_src/debugging.py
@@ -265,6 +265,10 @@ def debug_callback(callback: Callable[..., None], *args: Any,
 
 class _DebugPrintFormatChecker(string.Formatter):
 
+  def format_field(self, value, format_spec):
+    del value, format_spec
+    return ""  # No formatting is done.
+
   def check_unused_args(self, used_args, args, kwargs):
     unused_args = [arg for i, arg in enumerate(args) if i not in used_args]
     unused_kwargs = [k for k in kwargs if k not in used_args]
@@ -314,7 +318,7 @@ def debug_print(fmt: str, *args, ordered: bool = False, **kwargs) -> None:
     **kwargs: Additional keyword arguments to be formatted, as if passed to
       ``fmt.format``.
   """
-  # Check that we provide the correct arguments to be formatted
+  # Check that we provide the correct arguments to be formatted.
   formatter.format(fmt, *args, **kwargs)
 
   debug_callback(functools.partial(_format_print_callback, fmt), *args,


Test patch:
diff --git a/tests/debugging_primitives_test.py b/tests/debugging_primitives_test.py
index f5d7c4711..273c12f1b 100644
--- a/tests/debugging_primitives_test.py
+++ b/tests/debugging_primitives_test.py
@@ -106,6 +106,16 @@ class DebugPrintTest(jtu.JaxTestCase):
       jax.effects_barrier()
     self.assertEqual(output(), "x: 2\n")
 
+  def test_can_stage_out_debug_print_with_formatting(self):
+    @jax.jit
+    def f(x):
+      debug_print('x: {x:.2f}', x=x)
+
+    with jtu.capture_stdout() as output:
+      f(2)
+      jax.effects_barrier()
+    self.assertEqual(output(), "x: 2.00\n")
+
   @jtu.device_supports_buffer_donation()
   def test_can_stage_out_debug_print_with_donate_argnums(self):
     def f(x, y):