Make JAX tests that check for errors from dict key comparators in pytrees more relaxed, in preparation for openxla/xla#9529.

PiperOrigin-RevId: 610819296

Author: hawkinsp

tests/api_test.py

@@ -1329,7 +1329,9 @@ class E(enum.Enum):
     def f(d) -> float:
       return d[E.A]
 
-    with self.assertRaisesRegex(TypeError, "'<' not supported.*"):
+    with self.assertRaisesRegex(
+        (TypeError, ValueError),
+        "('<' not supported|Comparator raised exception).*"):
       f({E.A: 1.0, E.B: 2.0})
 
   def test_jit_static_argnums_requires_type_equality(self):

tests/tree_util_test.py

@@ -582,7 +582,9 @@ def testPickleRoundTrip(self, tree):
 
   def testDictKeysSortable(self):
     d = {"a": 1, 2: "b"}
-    with self.assertRaisesRegex(TypeError, "'<' not supported"):
+    with self.assertRaisesRegex(
+        (TypeError, ValueError),
+        "('<' not supported|Comparator raised exception).*"):
       _, _ = tree_util.tree_flatten(d)
 
   def testFlattenDictKeyOrder(self):