Don't include local_uint_constant_indices rewrite in JAX mode due to XLA bug

Author: ricardoV94

pytensor/compile/mode.py

@@ -451,7 +451,14 @@ def clone(self, link_kwargs=None, optimizer="", **kwargs):
     JAXLinker(),
     RewriteDatabaseQuery(
         include=["fast_run", "jax"],
-        exclude=["cxx_only", "BlasOpt", "fusion", "inplace"],
+        # TODO: "local_uint_constant_indices" can be reintroduced once https://github.com/google/jax/issues/16836 is fixed.
+        exclude=[
+            "cxx_only",
+            "BlasOpt",
+            "fusion",
+            "inplace",
+            "local_uint_constant_indices",
+        ],
     ),
 )
 NUMBA = Mode(