Don't include local_uint_constant_indices rewrite in JAX mode due to …

…XLA bug
pymc-devs · Jul 26, 2023 · 4459199 · 4459199
1 parent 14d2454
commit 4459199
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/pytensor/compile/mode.py b/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(