Bump jax commit

Author: wsmoses

src/enzyme_ad/jax/primitives.py

@@ -241,64 +241,8 @@ def stablehlo_inject(self):
     def ad_level(self):
         return self.passes.count("enzyme-wrap")
 
-
-DefaultCPPPipeline = OldXLAPipeline()  # NewXLAPipeline(None, True)
-DefaultJaXPipeline = JaXPipeline(
-    "inline{default-pipeline=canonicalize max-iterations=4},canonicalize,cse,enzyme-hlo-unroll,canonicalize,cse,enzyme-hlo-opt{passses=24575},cse"
-)
-
-
-def pass_pipeline(options):
-    if type(options) == type(""):
-        return options
-    else:
-        return
-
-
-def resource_dir():
-    import os
-
-    dn = os.path.dirname(enzyme_call.__file__)
-    res = os.path.join(dn, "..", "..", "clang", "staging")
-    return res
-
-
-def cflags():
-    import platform
-    import os
-
-    if platform.system() == "Darwin":
-        res = (
-            "-isysroot",
-            "/Library/Developer/CommandLineTools/SDKs/MacOSX.sdk",
-            "-isystem",
-            "/Applications/Xcode.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX.sdk/usr/include/c++/v1",
-            "-internal-isystem",
-            os.path.join(resource_dir(), "include"),
-            "-internal-externc-isystem",
-            "/Applications/Xcode.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX.sdk/usr/include",
-            "-internal-externc-isystem",
-            "/Applications/Xcode.app/Contents/Developer/Toolchains/XcodeDefault.xctoolchain/usr/include",
-            "-fgnuc-version=4.2.1",
-        )
-    else:
-        res = ()
-        if os.getenv("ENABLE_GDBLISTENER") is not None:
-            res = res + (
-                "-debug-info-kind=standalone",
-                "-dwarf-version=5",
-                "-debugger-tuning=gdb",
-            )
-    return res
-
-
-def optimize_module(mod, pipeline=None):
-    if pipeline is None:
-        pipeline = """
-            inline{default-pipeline=canonicalize max-iterations=4},
-            canonicalize,cse,
-            canonicalize,
-            enzyme-hlo-generate-td{
+def hlo_opts():
+    return """enzyme-hlo-generate-td{
             patterns=compare_op_canon<16>;
 broadcast_in_dim_op_canon<16>;
 convert_op_canon<16>;
@@ -452,6 +396,63 @@ def optimize_module(mod, pipeline=None):
             transform-interpreter,
             enzyme-hlo-remove-transform
         """
+
+DefaultCPPPipeline = OldXLAPipeline()  # NewXLAPipeline(None, True)
+DefaultJaXPipeline = JaXPipeline(
+    "inline{default-pipeline=canonicalize max-iterations=4},canonicalize,cse,enzyme-hlo-unroll,canonicalize,cse," + hlo_opts() + ", cse"
+)
+
+
+def pass_pipeline(options):
+    if type(options) == type(""):
+        return options
+    else:
+        return
+
+
+def resource_dir():
+    import os
+
+    dn = os.path.dirname(enzyme_call.__file__)
+    res = os.path.join(dn, "..", "..", "clang", "staging")
+    return res
+
+
+def cflags():
+    import platform
+    import os
+
+    if platform.system() == "Darwin":
+        res = (
+            "-isysroot",
+            "/Library/Developer/CommandLineTools/SDKs/MacOSX.sdk",
+            "-isystem",
+            "/Applications/Xcode.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX.sdk/usr/include/c++/v1",
+            "-internal-isystem",
+            os.path.join(resource_dir(), "include"),
+            "-internal-externc-isystem",
+            "/Applications/Xcode.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX.sdk/usr/include",
+            "-internal-externc-isystem",
+            "/Applications/Xcode.app/Contents/Developer/Toolchains/XcodeDefault.xctoolchain/usr/include",
+            "-fgnuc-version=4.2.1",
+        )
+    else:
+        res = ()
+        if os.getenv("ENABLE_GDBLISTENER") is not None:
+            res = res + (
+                "-debug-info-kind=standalone",
+                "-dwarf-version=5",
+                "-debugger-tuning=gdb",
+            )
+    return res
+
+
+def optimize_module(mod, pipeline=None):
+    if pipeline is None:
+        pipeline = """
+            inline{default-pipeline=canonicalize max-iterations=4},
+            canonicalize,cse,
+            canonicalize,""" + hlo_opts()
     enzyme_call.optimize_module(mod, pipeline)
     return
 
@@ -630,6 +631,7 @@ def maketup(ty):
     ty = ir.RankedTensorType(ty)
     tystr = ty.element_type.__str__()
     tystr = {
+        "i1": "bool",
         "bf16": "bfloat16",
         "f32": "float",
         "f64": "double",
@@ -1169,10 +1171,10 @@ def make_zero(tan, prim):
 
         outshapes = kwargs["out_shapes"]
         ret_act_tup = ",".join(["enzyme_dup"] * len(outshapes))
-        afterad = "arith-raise{stablehlo=true}, enzyme-hlo-opt, cse, canonicalize"
+        afterad = "arith-raise{stablehlo=true}, " + hlo_opts() + ", cse, canonicalize"
         newpasses = (
             "inline{default-pipeline=canonicalize max-iterations=4},"
-            + "enzyme-hlo-opt,cse,enzyme-wrap{infn=main outfn= retTys="
+            + hlo_opts() + ", cse,enzyme-wrap{infn=main outfn= retTys="
             + ret_act_tup
             + " argTys="
             + arg_act_tup
@@ -1401,7 +1403,7 @@ def enzyme_vjp(shadow_rets, *prim_args, **kwargs):
         newpasses = (
             prev_passes
             + ad_pass
-            + ",arith-raise{stablehlo=true},canonicalize, remove-unnecessary-enzyme-ops, enzyme-simplify-math, enzyme-hlo-opt, canonicalize, cse"
+            + ",arith-raise{stablehlo=true},canonicalize, remove-unnecessary-enzyme-ops, enzyme-simplify-math, " + hlo_opts() + ", canonicalize, cse"
             + post_passes
         )
 

workspace.bzl

@@ -1,5 +1,5 @@
-JAX_COMMIT = "8815b236b656f494171131301d1d81e84cf4c67c"
-JAX_SHA256 = "188787b8ec366dcda5805f24dcbfab7a349aa780498aeb9c3b728c9cec0a7e7d"
+JAX_COMMIT = "493698e6e053641aa8c51bca657cbd763a3ced19"
+JAX_SHA256 = ""
 
 ENZYME_COMMIT = "9acbc0a667ec8ae76407b5708758667a65ff15aa"
 ENZYME_SHA256 = "287143133ccf9501a02f1bdab351c34adcab3bbfc8648b180ebd79d0e058b3af"
@@ -32,4 +32,4 @@ XLA_PATCHES = [
     "find . -type f -name BUILD -exec sed -i.bak1 's/\\/\\/third_party\\/py\\/enzyme_ad\\/\\.\\.\\./public/g' {} +", 
     "find . -type f -name BUILD -exec sed -i.bak2 's/\\/\\/xla\\/mlir\\/memref:friends/\\/\\/visibility:public/g' {} +",
     "find xla/mlir -type f -name BUILD -exec sed -i.bak3 's/\\/\\/xla:internal/\\/\\/\\/\\/visibility:public/g' {} +"
-]
+]