Repository: jax-ml/jax
Issue category: Bug Report

Issue:
jax.lax.scan transforms dict keys to lower case when reporting mismatch in pytree structures
### Description

When `jax.lax.scan` reports an error due to mismatch in the pytree input/output structres, it transforms dict keys to lowercase.

Small repro:

```python
def f(loop_i, x):
  return {'T': jnp.array([0.5])}

init_val = {'t': jnp.array([1.0])}
jax.lax.fori_loop(0, 1, f, init_val)
```

Which leads to error:
```
TypeError: scan body function carry input and carry output must have the same pytree structure, but they differ:

The input carry component loop_carry[1] is a <class 'dict'> with pytree metadata ('t',) but the corresponding component of the carry output is a <class 'dict'> with pytree metadata ('t',), so the pytree node metadata does not match.
```

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

N/A


Hints:
I tracked it down to this line: https://github.com/jax-ml/jax/blob/2b06f93c703d062bebb6154a0dc030f2467c67cc/jax/_src/lax/control_flow/loops.py#L383

This causes all characters after the first within the error message to be converted to lowercase, which results in the error output you're seeing.

Patch:
diff --git a/jax/_src/lax/control_flow/loops.py b/jax/_src/lax/control_flow/loops.py
index b08819bcf545..5839ec3e50e4 100644
--- a/jax/_src/lax/control_flow/loops.py
+++ b/jax/_src/lax/control_flow/loops.py
@@ -347,6 +347,10 @@ def _get_states(attrs_tracked):
     vals.extend(leaves)
   return vals
 
+def _capitalize(s):
+  # s.capitalize() converts s[1:] to lowercase which we don't want.
+  return s[0].capitalize() + s[1:]
+
 def _check_carry_type(name, body_fun, in_carry, out_carry_tree, out_avals):
   try:
     sig = inspect.signature(body_fun)
@@ -380,7 +384,7 @@ def _check_carry_type(name, body_fun, in_carry, out_carry_tree, out_avals):
         # The trees may have different aux data but structures are the same.
         return
       if len(diffs) == 1:
-        differences = f'{diffs[0]}.\n'.capitalize()
+        differences = f'{_capitalize(diffs[0])}.\n'
       else:
         differences = ('\n'.join(f'  * {d};\n' for d in diffs[:-1])
                        + f'  * {diffs[-1]}.\n')
@@ -400,7 +404,7 @@ def _check_carry_type(name, body_fun, in_carry, out_carry_tree, out_avals):
       # The trees may have different aux data but structures are the same.
       return
     if len(diffs) == 1:
-      differences = f'{diffs[0]}.\n'.capitalize()
+      differences = f'{_capitalize(diffs[0])}.\n'
     else:
       differences = ('\n'.join(f'  * {d};\n' for d in diffs[:-1])
                      + f'  * {diffs[-1]}.\n')


Test patch:
diff --git a/tests/lax_control_flow_test.py b/tests/lax_control_flow_test.py
index 1f04a788914d..f323b035db6f 100644
--- a/tests/lax_control_flow_test.py
+++ b/tests/lax_control_flow_test.py
@@ -1897,6 +1897,16 @@ def testScanBodyOutputError(self):
         re.escape("scan body output must be a pair, got ShapedArray(float32[]).")):
       lax.scan(lambda c, x: np.float32(0.), 0, jnp.arange(5.))
 
+  def testScanMetadataError(self):
+    # Regression test for https://github.com/jax-ml/jax/issues/25507
+    def f(loop_i, x):
+      return {'T': jnp.array([0.5])}
+
+    init_val = {'t': jnp.array([1.0])}
+    msg = r".*with pytree metadata \('t',\).*with pytree metadata \('T',\)"
+    with self.assertRaisesRegex(TypeError, msg):
+      jax.lax.fori_loop(0, 1, f, init_val)
+
   def testScanBodyCarryPytreeMismatchErrors(self):
     with self.assertRaisesRegex(
         TypeError,