Repository: EnzymeAD/Enzyme-JAX (language: python)

Issue:
Make JAX import into WORKSPACE work.
The current import was non-functional, and that only remained unnoticed
because there is no BUILD dependency to JAX, so JAX isn't actually
built. (The CI tests install JAX through pip, so the JAX dependency
isn't needed there.)

The fix consists in moving the configuration from JAX's WORKSPACE file
(which is emptied when imported via `http_archive`) to the WORKSPACE
file of this project (plus fixing how its dependency `ducc` is loaded).
Now adding `@jax//jax` to some of the build targets actually (and
successfully) builds JAX.


Discussion / hints:

Fix patch:
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index 7febf1f5..506e54d5 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -35,13 +35,18 @@ jobs:
           path: "~/.cache/bazel"
           key: bazel-${{ matrix.os }}
       - run: find ~/.cache/bazel ~/.cache/bazelisk -iname "*.whl" -exec rm {} \;
-      - run: bazel build :enzyme_jax @llvm-project//llvm:FileCheck
+      - run: |
+          bazel build :enzyme_jax @llvm-project//llvm:FileCheck
+          bazel cquery "allpaths(//enzyme_jax:enzyme_call,@xla//xla/stream_executor:executor_cache)" --notool_deps
+          bazel --version
+          nm -C $(find bazel-out/ -name enzyme_call.so -type f) | grep ExecutorCache::
       - run: cp bazel-bin/*.whl .
       
       - name: test
         run: |
           python3 -m pip install --user --force-reinstall "jax[cpu]" *.whl
           cd test
+          nm -C $(python3 -c "from enzyme_jax import enzyme_call; print(enzyme_call.__file__)") | grep ExecutorCache::
           python3 test.py
           cd lit_tests
           lit . --verbose
diff --git a/BUILD b/BUILD
index 5029ef3e..e0eb3862 100644
--- a/BUILD
+++ b/BUILD
@@ -1,8 +1,6 @@
 
 licenses(["notice"])
 
-load("@tsl//tsl:tsl.bzl", _if_windows = "if_windows", pybind_extension = "tsl_pybind_extension_opensource")
-
 package(
     default_applicable_licenses = [],
     default_visibility = ["//:__subpackages__"],
@@ -30,18 +28,17 @@ cc_library(
 )
 
 load("@rules_python//python:packaging.bzl", "py_wheel")
-
 load(":package.bzl", "py_package")
 
 py_package(
     name = "enzyme_jax_data",
-    # Only include these Python packages.
-    packages = ["@//enzyme_jax:enzyme_call.so", "@llvm-project//clang:builtin_headers_gen"],
     deps = [
-        "//enzyme_jax:enzyme_call",
+        "//enzyme_jax:enzyme_call.so",
         "@llvm-project//clang:builtin_headers_gen",
     ],
-    prefix = "enzyme_jax/",
+    # Only include these Python packages.
+    packages = ["@//enzyme_jax:enzyme_call.so", "@llvm-project//clang:builtin_headers_gen"],
+    prefix = "enzyme_jax/"
 )
 
 py_wheel(
diff --git a/WORKSPACE b/WORKSPACE
index 02206ff9..92cb3e42 100644
--- a/WORKSPACE
+++ b/WORKSPACE
@@ -1,7 +1,21 @@
 load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 
-LLVM_COMMIT = "aa495214b39d475bab24b468de7a7c676ce9e366"
-LLVM_SHA256 = "73cb1e91901d975bf4c97f1ea7000dd1554ad77f704a2d899498866a67471444"
+http_archive(
+    name = "rules_cc",
+    sha256 = "85723d827f080c5e927334f1fb18a294c0b3f94fee6d6b45945f5cdae6ea0fd4",
+    strip_prefix = "rules_cc-c8c38f8c710cbbf834283e4777916b68261b359c",
+    urls = [
+        "https://github.com/bazelbuild/rules_cc/archive/c8c38f8c710cbbf834283e4777916b68261b359c.tar.gz",
+    ],
+)
+
+load("@rules_cc//cc:repositories.bzl", "rules_cc_dependencies")
+
+rules_cc_dependencies()
+
+LLVM_COMMIT = "5e5a22caf88ac1ccfa8dc5720295fdeba0ad9372"
+LLVM_SHA256 = ""
+LLVM_TARGETS = ["X86", "AArch64", "AMDGPU", "NVPTX"]
 
 http_archive(
     name = "llvm-raw",
@@ -14,10 +28,10 @@ http_archive(
 )
 
 load("@llvm-raw//utils/bazel:configure.bzl", "llvm_configure")
-llvm_configure(name = "llvm-project", targets = ["X86", "AArch64", "AMDGPU", "ARM", "NVPTX"])
+llvm_configure(name = "llvm-project", targets = LLVM_TARGETS)
 
-XLA_COMMIT = "7423e38a383ccd25fb144db298257a6b6dd8cc4d"
-XLA_SHA256 = "e7ea840e4a58a91bdb5dbcaee71669dd799cecc8d2084106389699619fe76129"
+XLA_COMMIT = "fecd5e7e9f00f4a197ad54206f2bc0ca1058c858"
+XLA_SHA256 = ""
 
 http_archive(
     name = "xla",
@@ -28,9 +42,8 @@ http_archive(
     patches = ["//:patches/xla.patch"],
 )
 
-
-PYRULES_COMMIT = "693a1587baf055979493565933f8f40225c00c6d"
-PYRULES_SHA256 = "c493a9506b5e1ea99e3c22fb15e672cdd2a6fa19fd7c627ec5d485aced23e50a"
+PYRULES_COMMIT = "fe33a4582c37499f3caeb49a07a78fc7948a8949"
+PYRULES_SHA256 = "cfa6957832ae0e0c7ee2ccf455a888a291e8419ed8faf45f4420dd7414d5dd96"
 
 http_archive(
     name = "rules_python",
@@ -39,8 +52,16 @@ http_archive(
     urls = ["https://github.com/bazelbuild/rules_python/archive/{commit}.tar.gz".format(commit = PYRULES_COMMIT)]
 )
 
-ENZYME_COMMIT = "bcd061afc6260d2266ca9a8489830c36a4ceefe6"
-ENZYME_SHA256 = "f215f6654b000a7eed387d89fe51561382dc7a86e1ed83941399335f819c1f66"
+load("@rules_python//python:repositories.bzl", "py_repositories")
+
+py_repositories()
+
+load("@rules_python//python/pip_install:repositories.bzl", "pip_install_dependencies")
+
+pip_install_dependencies()
+
+ENZYME_COMMIT = "77b4fff47701a240b537a93a2e722626f7421342"
+ENZYME_SHA256 = ""
 
 http_archive(
     name = "enzyme",
@@ -49,6 +70,21 @@ http_archive(
     urls = ["https://github.com/EnzymeAD/Enzyme/archive/{commit}.tar.gz".format(commit = ENZYME_COMMIT)],
 )
 
+JAX_COMMIT = "32a317f7a43440800e1e39e00ed5f2980e088ab1"
+JAX_SHA256 = "6e2147be7360a5c0672b6ba0d654cdb2ac96113b63ef457dfdc76cd50fe69ff1"
+
+http_archive(
+    name = "jax",
+    sha256 = JAX_SHA256,
+    strip_prefix = "jax-" + JAX_COMMIT,
+    urls = ["https://github.com/google/jax/archive/{commit}.tar.gz".format(commit = JAX_COMMIT)],
+    patch_args = ["-p1"],
+    patches = ["//:patches/jax.patch"],
+)
+
+load("@jax//third_party/xla:workspace.bzl", jax_xla_workspace = "repo")
+jax_xla_workspace()
+
 load("@xla//:workspace4.bzl", "xla_workspace4")
 xla_workspace4()
 
@@ -64,15 +100,11 @@ xla_workspace1()
 load("@xla//:workspace0.bzl", "xla_workspace0")
 xla_workspace0()
 
+load("@jax//third_party/flatbuffers:workspace.bzl", flatbuffers = "repo")
+flatbuffers()
 
-JAX_COMMIT = "21fc6e0229e0f5f1cb5f1f69d2c3daa2e5c2ca11"
-JAX_SHA256 = ""
+load("@jax//third_party/robin_map:workspace.bzl", robin_map = "repo")
+robin_map()
 
-http_archive(
-    name = "jax",
-    sha256 = JAX_SHA256,
-    strip_prefix = "jax-" + PYRULES_COMMIT,
-    urls = ["https://github.com/google/jax/archive/{commit}.tar.gz".format(commit = JAX_COMMIT)],
-    patch_args = ["-p1"],
-    patches = ["//:patches/jax_workspace.patch"]
-)
+load("@jax//third_party/nanobind:workspace.bzl", nanobind = "repo")
+nanobind()
diff --git a/clang_compile.cc b/clang_compile.cc
index 167b15df..fce4dfcc 100644
--- a/clang_compile.cc
+++ b/clang_compile.cc
@@ -147,45 +147,10 @@ template <class T> class ptr_wrapper
 PYBIND11_DECLARE_HOLDER_TYPE(T, ptr_wrapper<T>, true);
 */
 
-static TargetLibraryInfoImpl *createTLII(llvm::Triple &&TargetTriple,
-                                         const CodeGenOptions &CodeGenOpts) {
-  TargetLibraryInfoImpl *TLII = new TargetLibraryInfoImpl(TargetTriple);
- 
-  switch (CodeGenOpts.getVecLib()) {
-  case CodeGenOptions::Accelerate:
-    TLII->addVectorizableFunctionsFromVecLib(TargetLibraryInfoImpl::Accelerate,
-                                             TargetTriple);
-    break;
-  case CodeGenOptions::LIBMVEC:
-    TLII->addVectorizableFunctionsFromVecLib(TargetLibraryInfoImpl::LIBMVEC_X86,
-                                             TargetTriple);
-    break;
-  case CodeGenOptions::MASSV:
-    TLII->addVectorizableFunctionsFromVecLib(TargetLibraryInfoImpl::MASSV,
-                                             TargetTriple);
-    break;
-  case CodeGenOptions::SVML:
-    TLII->addVectorizableFunctionsFromVecLib(TargetLibraryInfoImpl::SVML,
-                                             TargetTriple);
-    break;
-  case CodeGenOptions::SLEEF:
-    TLII->addVectorizableFunctionsFromVecLib(TargetLibraryInfoImpl::SLEEFGNUABI,
-                                             TargetTriple);
-    break;
-  case CodeGenOptions::Darwin_libsystem_m:
-    TLII->addVectorizableFunctionsFromVecLib(
-        TargetLibraryInfoImpl::DarwinLibSystemM, TargetTriple);
-    break;
-  default:
-    break;
-  }
-  return TLII;
-}
-
 // Returns the TargetMachine instance or zero if no triple is provided.
 static TargetMachine* GetTargetMachine(llvm::Triple TheTriple, StringRef CPUStr,
                                        StringRef FeaturesStr,
-                                       const llvm::TargetOptions &Options, CodeGenOpt::Level level) {
+                                       const llvm::TargetOptions &Options, CodeGenOptLevel level) {
   std::string Error;
   const Target *TheTarget =
       TargetRegistry::lookupTarget(codegen::getMArch(), TheTriple, Error);
@@ -548,12 +513,13 @@ struct tensor<T, n0, N...>
 
   // Register the target library analysis directly and give it a customized
   // preset TLI.
+  llvm::Triple triple(mod->getTargetTriple());
   std::unique_ptr<TargetLibraryInfoImpl> TLII(
-      createTLII(llvm::Triple(mod->getTargetTriple()), Clang->getCodeGenOpts()));
+      llvm::driver::createTLII(triple, Clang->getCodeGenOpts().getVecLib()));
   FAM.registerPass([&] { return TargetLibraryAnalysis(*TLII); });
 
 
-  auto level = CodeGenOpt::Level::Aggressive; //OptimizationLevel::O3;
+  auto level = CodeGenOptLevel::Aggressive; //OptimizationLevel::O3;
 
   Triple ModuleTriple(mod->getTargetTriple());
   std::string CPUStr, FeaturesStr;
diff --git a/enzyme_jax/BUILD b/enzyme_jax/BUILD
index d1a28780..a968e989 100644
--- a/enzyme_jax/BUILD
+++ b/enzyme_jax/BUILD
@@ -1,5 +1,5 @@
 
-load("@tsl//tsl:tsl.bzl", _if_windows = "if_windows", pybind_extension = "tsl_pybind_extension_opensource")
+load("@pybind11_bazel//:build_defs.bzl", "pybind_extension", "pybind_library")
 
 py_library(
     name = "enzyme_jax_internal",
@@ -7,7 +7,7 @@ py_library(
     visibility = ["//visibility:public"]
 )
 
-cc_library(
+pybind_library(
     name = "compile_with_xla",
     srcs = ["compile_with_xla.cc"],
     deps = [
@@ -20,8 +20,6 @@ cc_library(
         "@tsl//tsl/profiler/backends/cpu:annotation_stack_impl",
         "@tsl//tsl/profiler/backends/cpu:traceme_recorder_impl",
         "@tsl//tsl/profiler/utils:time_utils_impl",
-        "@tsl//tsl/protobuf:autotuning_proto_cc_impl",
-        "@tsl//tsl/protobuf:dnn_proto_cc",
         "@tsl//tsl/protobuf:dnn_proto_cc_impl",
         "@tsl//tsl/protobuf:histogram_proto_cc",
         "@tsl//tsl/protobuf:histogram_proto_cc_impl",
@@ -29,6 +27,8 @@ cc_library(
         "@tsl//tsl/util:determinism",
 
         # This is similar to xla_binary rule and is needed to make XLA client compile.
+        "@xla//xla:autotuning_proto_cc",
+        "@xla//xla:autotuning_proto_cc_impl",
         "@xla//xla:autotune_results_proto_cc",
         "@xla//xla:autotune_results_proto_cc_impl",
         "@xla//xla/client",
@@ -36,10 +36,14 @@ cc_library(
         "@xla//xla/service/cpu:cpu_executable",
         "@xla//xla/service/gpu:backend_configs_cc",
         "@xla//xla/service/gpu:backend_configs_cc_impl",
+        "@xla//xla/service:buffer_assignment_proto_cc",
+        "@xla//xla/service:buffer_assignment_proto_cc_impl",
         "@xla//xla/service:hlo_proto_cc",
         "@xla//xla/service:hlo_proto_cc_impl",
-        "@xla//xla/service:memory_space_assignment_proto_cc_impl",
-        "@xla//xla/stream_executor:dnn_proto_cc_impl",
+        "@xla//xla/service/memory_space_assignment:memory_space_assignment_proto_cc_impl",
+        "@xla//xla/stream_executor:stream_executor_impl",
+        "@xla//xla/stream_executor:device_description_proto_cc",
+        "@xla//xla/stream_executor:device_description_proto_cc_impl",
         "@xla//xla:xla_data_proto_cc",
         "@xla//xla:xla_data_proto_cc_impl",
         "@xla//xla:xla_proto_cc",
@@ -66,7 +70,6 @@ pybind_extension(
     name = "enzyme_call",
     srcs = ["enzyme_call.cc"],
     deps = [
-        "@pybind11",
         "@llvm-project//llvm:Support",
         "@llvm-project//llvm:OrcJIT",
         "//:clang_compile",
diff --git a/enzyme_jax/primitives.py b/enzyme_jax/primitives.py
index 77658272..31ebfad7 100644
--- a/enzyme_jax/primitives.py
+++ b/enzyme_jax/primitives.py
@@ -16,8 +16,7 @@
 
 import jax.numpy as jnp
 
-#from build import enzyme_call
-from enzyme_jax import enzyme_call
+from . import enzyme_call
 
 LANG_CPP = enzyme_call.Language.CPP
 LANG_LLVM = enzyme_call.Language.LLVM
diff --git a/patches/jax.patch b/patches/jax.patch
new file mode 100644
index 00000000..1fb3346b
--- /dev/null
+++ b/patches/jax.patch
@@ -0,0 +1,11 @@
+--- a/jaxlib/cpu/BUILD
++++ a/jaxlib/cpu/BUILD
+@@ -79,7 +79,7 @@ cc_library(
+         ":ducc_fft_flatbuffers_cc",
+         "@xla//xla/service:custom_call_status",
+         "@com_github_google_flatbuffers//:flatbuffers",
+-        "@ducc//:fft",
++        "@ducc//:fft_wrapper",
+     ],
+ )
+ 
diff --git a/patches/jax_workspace.patch b/patches/jax_workspace.patch
deleted file mode 100644
index ade67761..00000000
--- a/patches/jax_workspace.patch
+++ /dev/null
@@ -1,45 +0,0 @@
-diff --git a/WORKSPACE b/WORKSPACE
-index 6250624e9..33f473cb3 100644
---- a/WORKSPACE
-+++ b/WORKSPACE
-@@ -1,17 +1,33 @@
--load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
-+#load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
- 
- # To update XLA to a new revision,
- # a) update URL and strip_prefix to the new git commit hash
- # b) get the sha256 hash of the commit by running:
- #    curl -L https://github.com/openxla/xla/archive/<git hash>.tar.gz | sha256sum
- #    and update the sha256 with the result.
--http_archive(
-+#http_archive(
-+#    name = "xla",
-+#    sha256 = "14bbdd796211e45a7a3148bd5e43523da4556ff19cf5461588b79650b55b9302",
-+#    strip_prefix = "xla-c1e4a16e77a7ba2000003ccade3ffba3749ada35",
-+#    urls = [
-+#        "https://github.com/openxla/xla/archive/c1e4a16e77a7ba2000003ccade3ffba3749ada35.tar.gz",
-+#    ],
-+#)
-+
-+new_local_repository(
-+    name = "llvm-raw",
-+    build_file = "llvm_dummy.BUILD",
-+    path = "../llvm-project",
-+)
-+
-+load("@llvm-raw//utils/bazel:configure.bzl", "llvm_configure", "llvm_disable_optional_support_deps")
-+
-+llvm_disable_optional_support_deps()
-+llvm_configure(name = "llvm-project", targets = ["X86"])
-+
-+local_repository(
-     name = "xla",
--    sha256 = "14bbdd796211e45a7a3148bd5e43523da4556ff19cf5461588b79650b55b9302",
--    strip_prefix = "xla-c1e4a16e77a7ba2000003ccade3ffba3749ada35",
--    urls = [
--        "https://github.com/openxla/xla/archive/c1e4a16e77a7ba2000003ccade3ffba3749ada35.tar.gz",
--    ],
-+    path = "../xla",
- )
- 
- # For development, one often wants to make changes to the TF repository as well
diff --git a/patches/xla.patch b/patches/xla.patch
index 087ad4be..5e799e11 100644
--- a/patches/xla.patch
+++ b/patches/xla.patch
@@ -1,197 +1,18 @@
-diff '--color=auto' -ru a/xla/BUILD b/xla/BUILD
---- a/xla/BUILD	2023-06-12 20:50:10.000000000 +0000
-+++ b/xla/BUILD	2023-07-11 09:42:42.089247113 +0000
-@@ -13,7 +13,7 @@
- 
- package(
-     # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
--    default_visibility = ["//xla:internal"],
-+    default_visibility = ["//visibility:public"],
-     licenses = ["notice"],
- )
- 
-diff '--color=auto' -ru a/xla/mlir/backends/cpu/transforms/legalize_i1_vector_transfers.cc b/xla/mlir/backends/cpu/transforms/legalize_i1_vector_transfers.cc
---- a/xla/mlir/backends/cpu/transforms/legalize_i1_vector_transfers.cc	2023-06-12 20:50:10.000000000 +0000
-+++ b/xla/mlir/backends/cpu/transforms/legalize_i1_vector_transfers.cc	2023-07-10 15:50:07.787003636 +0000
-@@ -89,7 +89,7 @@
-         DenseElementsAttr::get(new_read.getType(), b.getI8IntegerAttr(0)));
-     auto result =
-         b.create<arith::CmpIOp>(arith::CmpIPredicate::ne, new_read, zero);
--    rewriter.replaceOp(op, {result});
-+    rewriter.replaceOp(op, result);
-     return success();
-   };
- };
-diff '--color=auto' -ru a/xla/mlir_hlo/mhlo/IR/hlo_ops.cc b/xla/mlir_hlo/mhlo/IR/hlo_ops.cc
---- a/xla/mlir_hlo/mhlo/IR/hlo_ops.cc	2023-06-12 20:50:10.000000000 +0000
-+++ b/xla/mlir_hlo/mhlo/IR/hlo_ops.cc	2023-07-10 15:55:08.809414041 +0000
-@@ -3005,7 +3005,7 @@
-         rewriter.getI64TensorAttr(SmallVector<int64_t, 4>(inputRank, 1));
-     auto result = rewriter.create<SliceOp>(loc, input, sliceStartIndices,
-                                            sliceLimits, sliceStrides);
--    rewriter.replaceOp(dynamicSlice, {result});
-+    rewriter.replaceOp(dynamicSlice, result);
-     return success();
-   }
- };
-diff '--color=auto' -ru a/xla/mlir_hlo/mhlo/transforms/chlo_legalize_to_hlo/chlo_legalize_to_hlo.cc b/xla/mlir_hlo/mhlo/transforms/chlo_legalize_to_hlo/chlo_legalize_to_hlo.cc
---- a/xla/mlir_hlo/mhlo/transforms/chlo_legalize_to_hlo/chlo_legalize_to_hlo.cc	2023-06-12 20:50:10.000000000 +0000
-+++ b/xla/mlir_hlo/mhlo/transforms/chlo_legalize_to_hlo/chlo_legalize_to_hlo.cc	2023-07-10 15:52:49.617201134 +0000
-@@ -1797,8 +1797,8 @@
-     }
- 
-     rewriter.replaceOp(op,
--                       {Adaptor::createOp(op, op.getResult().getType(),
--                                          adaptor.getOperands(), rewriter)});
-+                       Adaptor::createOp(op, op.getResult().getType(),
-+                                          adaptor.getOperands(), rewriter));
-     return success();
-   }
- };
-diff '--color=auto' -ru a/xla/python/BUILD b/xla/python/BUILD
---- a/xla/python/BUILD	2023-06-12 20:50:10.000000000 +0000
-+++ b/xla/python/BUILD	2023-07-11 09:44:12.489187435 +0000
-@@ -21,7 +21,7 @@
- 
- package(
-     # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
--    default_visibility = [":friends"],
-+    default_visibility = ["//visibility:public"],
-     licenses = ["notice"],
- )
- 
-@@ -70,7 +70,7 @@
-     testonly = 1,
-     srcs = ["xla_client_test.py"],
-     srcs_version = "PY3",
--    visibility = [":friends"],
-+    visibility = ["//visibility:public"],
-     deps = [
-         ":xla_client",
-         "@absl_py//absl/flags",
-@@ -205,7 +205,7 @@
-         "-fno-strict-aliasing",
-     ],
-     features = ["-use_header_modules"],
--    visibility = [":friends"],
-+    visibility = ["//visibility:public"],
-     deps = [
-         ":exceptions",
-         "//xla:literal",
-@@ -232,7 +232,7 @@
-         "-fno-strict-aliasing",
-     ],
-     features = ["-use_header_modules"],
--    visibility = [":friends"],
-+    visibility = ["//visibility:public"],
-     deps = [
-         "@com_google_absl//absl/base:core_headers",
-         "@com_google_absl//absl/container:inlined_vector",
-@@ -268,7 +268,7 @@
-         "-fno-strict-aliasing",
+--- a/xla/BUILD
++++ b/xla/BUILD
+@@ -29,6 +29,7 @@ package_group(
+         "//third_party/py/tpu_graphs/...",
+         "//tensorflow/compiler/...",
+         "//tensorflow/python/tpu/...",
++        "public",
      ],
-     features = ["-use_header_modules"],
--    visibility = [":friends"],
-+    visibility = ["//visibility:public"],
-     deps = [
-         ":exceptions",
-         ":python_ref_manager",
-@@ -435,7 +435,7 @@
-         "-fno-strict-aliasing",
-     ],
-     features = ["-use_header_modules"],
--    visibility = [":friends"],  # For the functions to access C++ flags/thread-local variables
-+    visibility = ["//visibility:public"],  # For the functions to access C++ flags/thread-local variables
-     deps = [
-         ":exceptions",
-         ":py_client",
-@@ -692,7 +692,7 @@
-         "-fno-strict-aliasing",
-     ],
-     features = ["-use_header_modules"],
--    visibility = [":friends"],
-+    visibility = ["//visibility:public"],
-     deps = [
-         ":exceptions",
-         ":pytree_proto_cc",
-@@ -801,7 +801,7 @@
-         "-fno-strict-aliasing",
-     ],
-     features = ["-use_header_modules"],
--    visibility = [":friends"],
-+    visibility = ["//visibility:public"],
-     deps = [
-         ":status_casters",
-         "//xla:status",
-diff '--color=auto' -ru a/xla/runtime/BUILD b/xla/runtime/BUILD
---- a/xla/runtime/BUILD	2023-06-12 20:50:10.000000000 +0000
-+++ b/xla/runtime/BUILD	2023-07-12 14:00:16.882174573 +0000
-@@ -282,8 +282,6 @@
-         "@llvm-project//llvm:OrcJIT",
-         "@llvm-project//llvm:Support",
-         "@llvm-project//llvm:TransformUtils",
--        "@llvm-project//llvm:X86AsmParser",
--        "@llvm-project//llvm:X86CodeGen",
-     ] + select({
-         "@tsl//tsl:arm_any": [
-             "@llvm-project//llvm:AArch64AsmParser",  # fixdeps: keep
-@@ -298,6 +296,8 @@
-             "@llvm-project//llvm:AArch64CodeGen",  # fixdeps: keep
-         ],
-         "//conditions:default": [
-+            "@llvm-project//llvm:X86AsmParser",
-+            "@llvm-project//llvm:X86CodeGen",
-         ],
-     }) + if_llvm_system_z_available([
-         "@llvm-project//llvm:SystemZAsmParser",  # fixdeps: keep
-diff '--color=auto' -ru a/xla/service/BUILD b/xla/service/BUILD
---- a/xla/service/BUILD	2023-06-12 20:50:10.000000000 +0000
-+++ b/xla/service/BUILD	2023-07-10 15:36:23.550756903 +0000
-@@ -33,7 +33,7 @@
- 
- package(
-     # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
--    default_visibility = [":friends"],
-+    default_visibility = ["//visibility:public"],
-     licenses = ["notice"],
  )
  
-diff '--color=auto' -ru a/xla/service/cpu/BUILD b/xla/service/cpu/BUILD
---- a/xla/service/cpu/BUILD	2023-06-12 20:50:10.000000000 +0000
-+++ b/xla/service/cpu/BUILD	2023-07-10 14:56:23.459895361 +0000
-@@ -325,7 +325,6 @@
-         "@llvm-project//llvm:Support",
-         "@llvm-project//llvm:Target",
-         "@llvm-project//llvm:TargetParser",
--        "@llvm-project//llvm:X86CodeGen",  # fixdeps: keep
-         "@llvm-project//mlir:AffineDialect",
-         "@llvm-project//mlir:AffineToStandard",
-         "@llvm-project//mlir:ArithDialect",
-@@ -361,6 +360,7 @@
-             "@llvm-project//llvm:AArch64CodeGen",  # fixdeps: keep
-         ],
-         "//conditions:default": [
-+            "@llvm-project//llvm:X86CodeGen",  # fixdeps: keep
-         ],
-     }) + if_llvm_system_z_available([
-         "@llvm-project//llvm:SystemZCodeGen",  # fixdeps: keep
-@@ -559,6 +559,7 @@
-         "@tsl//tsl/platform:env",
-         "@tsl//tsl/platform:logging",
-     ],
+@@ -1165,6 +1166,7 @@ tf_proto_library(
+     srcs = ["autotuning.proto"],
+     make_default_target_header_only = True,
+     protodeps = ["@tsl//tsl/protobuf:dnn_proto"],
 +    visibility = ["//visibility:public"],
  )
  
  cc_library(
-diff '--color=auto' -ru a/xla/service/gpu/BUILD b/xla/service/gpu/BUILD
---- a/xla/service/gpu/BUILD	2023-06-12 20:50:10.000000000 +0000
-+++ b/xla/service/gpu/BUILD	2023-07-10 15:28:03.306876220 +0000
-@@ -30,7 +30,7 @@
- 
- package(
-     # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
--    default_visibility = [":friends"],
-+    default_visibility = ["//visibility:public"],
-     licenses = ["notice"],
- )
- 


Test patch:
diff --git a/test/lit_tests/ir.pyt b/test/lit_tests/ir.pyt
index 4ea1b892..e3198cc9 100644
--- a/test/lit_tests/ir.pyt
+++ b/test/lit_tests/ir.pyt
@@ -34,7 +34,7 @@ def fwdmode(a, b, c, d):
 print(fwdmode.lower(ones, twos, ones, twos).compiler_ir(dialect="mhlo"))
 
 # CHECK: module @jit_fwdmode attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
-# CHECK-NEXT:   func.func public @main(%arg0: tensor<2x3xf32> {jax.arg_info = "a", mhlo.sharding = "{replicated}"}, %arg1: tensor<5x7xf32> {jax.arg_info = "b", mhlo.sharding = "{replicated}"}, %arg2: tensor<2x3xf32> {jax.arg_info = "c", mhlo.sharding = "{replicated}"}, %arg3: tensor<5x7xf32> {jax.arg_info = "d", mhlo.sharding = "{replicated}"}) -> (tensor<6x9xf32> {jax.result_info = "[0][0]"}, tensor<4x6xf32> {jax.result_info = "[0][1]"}, tensor<6x9xf32> {jax.result_info = "[1][0]"}, tensor<4x6xf32> {jax.result_info = "[1][1]"}) {
+# CHECK-NEXT:   func.func public @main(%arg0: tensor<2x3xf32> {mhlo.sharding = "{replicated}"}, %arg1: tensor<5x7xf32> {mhlo.sharding = "{replicated}"}, %arg2: tensor<2x3xf32> {mhlo.sharding = "{replicated}"}, %arg3: tensor<5x7xf32> {mhlo.sharding = "{replicated}"}) -> (tensor<6x9xf32> {jax.result_info = "[0][0]"}, tensor<4x6xf32> {jax.result_info = "[0][1]"}, tensor<6x9xf32> {jax.result_info = "[1][0]"}, tensor<4x6xf32> {jax.result_info = "[1][1]"}) {
 # CHECK-NEXT:     %0 = mhlo.constant dense<1> : tensor<1xi64>
 # CHECK-NEXT:     %1:4 = mhlo.custom_call @jaxzyme.fwd(%0, %arg0, %arg2, %arg1, %arg3) {backend_config = ""} : (tensor<1xi64>, tensor<2x3xf32>, tensor<2x3xf32>, tensor<5x7xf32>, tensor<5x7xf32>) -> (tensor<6x9xf32>, tensor<6x9xf32>, tensor<4x6xf32>, tensor<4x6xf32>)
 # CHECK-NEXT:     return %1#0, %1#2, %1#1, %1#3 : tensor<6x9xf32>, tensor<4x6xf32>, tensor<6x9xf32>, tensor<4x6xf32>
@@ -48,7 +48,7 @@ def f(a, b):
 print(f.lower(ones, twos).compiler_ir(dialect="mhlo"))
 
 # CHECK: module @jit_f attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
-# CHECK-NEXT:   func.func public @main(%arg0: tensor<2x3xf32> {jax.arg_info = "a", mhlo.sharding = "{replicated}"}, %arg1: tensor<5x7xf32> {jax.arg_info = "b", mhlo.sharding = "{replicated}"}) -> (tensor<6x9xf32> {jax.result_info = "[0][0]"}, tensor<4x6xf32> {jax.result_info = "[0][1]"}, tensor<16xi8> {jax.result_info = "[1][<flat index 0>][0][<flat index 0>][0][0]"}, tensor<2x3xf32> {jax.result_info = "[1][<flat index 0>][0][<flat index 0>][0][1]"}, tensor<5x7xf32> {jax.result_info = "[1][<flat index 0>][0][<flat index 0>][0][2]"}) {
+# CHECK-NEXT:   func.func public @main(%arg0: tensor<2x3xf32> {mhlo.sharding = "{replicated}"}, %arg1: tensor<5x7xf32> {mhlo.sharding = "{replicated}"}) -> (tensor<6x9xf32> {jax.result_info = "[0][0]"}, tensor<4x6xf32> {jax.result_info = "[0][1]"}, tensor<16xi8> {jax.result_info = "[1][<flat index 0>][0][<flat index 0>][0][0]"}, tensor<2x3xf32> {jax.result_info = "[1][<flat index 0>][0][<flat index 0>][0][1]"}, tensor<5x7xf32> {jax.result_info = "[1][<flat index 0>][0][<flat index 0>][0][2]"}) {
 # CHECK-NEXT:     %0 = mhlo.constant dense<2> : tensor<1xi64>
 # CHECK-NEXT:     %1:3 = mhlo.custom_call @jaxzyme.aug(%0, %arg0, %arg1) {backend_config = ""} : (tensor<1xi64>, tensor<2x3xf32>, tensor<5x7xf32>) -> (tensor<6x9xf32>, tensor<4x6xf32>, tensor<16xi8>)
 # CHECK-NEXT:     return %1#0, %1#1, %1#2, %arg0, %arg1 : tensor<6x9xf32>, tensor<4x6xf32>, tensor<16xi8>, tensor<2x3xf32>, tensor<5x7xf32>
@@ -64,7 +64,7 @@ def g(a, b, x, y):
     return primals, f_vjp((x, y))
 
 # CHECK: module @jit_g attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
-# CHECK-NEXT:   func.func public @main(%arg0: tensor<2x3xf32> {jax.arg_info = "a", mhlo.sharding = "{replicated}"}, %arg1: tensor<5x7xf32> {jax.arg_info = "b", mhlo.sharding = "{replicated}"}, %arg2: tensor<6x9xf32> {jax.arg_info = "x", mhlo.sharding = "{replicated}"}, %arg3: tensor<4x6xf32> {jax.arg_info = "y", mhlo.sharding = "{replicated}"}) -> (tensor<6x9xf32> {jax.result_info = "[0][0]"}, tensor<4x6xf32> {jax.result_info = "[0][1]"}, tensor<2x3xf32> {jax.result_info = "[1][0]"}, tensor<5x7xf32> {jax.result_info = "[1][1]"}) {
+# CHECK-NEXT:   func.func public @main(%arg0: tensor<2x3xf32> {mhlo.sharding = "{replicated}"}, %arg1: tensor<5x7xf32> {mhlo.sharding = "{replicated}"}, %arg2: tensor<6x9xf32> {mhlo.sharding = "{replicated}"}, %arg3: tensor<4x6xf32> {mhlo.sharding = "{replicated}"}) -> (tensor<6x9xf32> {jax.result_info = "[0][0]"}, tensor<4x6xf32> {jax.result_info = "[0][1]"}, tensor<2x3xf32> {jax.result_info = "[1][0]"}, tensor<5x7xf32> {jax.result_info = "[1][1]"}) {
 # CHECK-NEXT:     %0 = mhlo.constant dense<3> : tensor<1xi64>
 # CHECK-NEXT:     %1:3 = mhlo.custom_call @jaxzyme.aug(%0, %arg0, %arg1) {backend_config = ""} : (tensor<1xi64>, tensor<2x3xf32>, tensor<5x7xf32>) -> (tensor<6x9xf32>, tensor<4x6xf32>, tensor<16xi8>)
 # CHECK-NEXT:     %2 = mhlo.constant dense<4> : tensor<1xi64>