Merge pull request jax-ml#25511 from jakevdp:fix-scan-err

PiperOrigin-RevId: 706780339

Author: Google-ML-Automation

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')

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,