diff --git a/deps/ReactantExtra/API.cpp b/deps/ReactantExtra/API.cpp
index 83929dc21e..18e8a5ce2e 100644
--- a/deps/ReactantExtra/API.cpp
+++ b/deps/ReactantExtra/API.cpp
@@ -94,6 +94,7 @@
 #include "xla/python/ifrt/topology.h"
 #include "xla/python/ifrt/tuple.h"
 #include "xla/python/ifrt/value.h"
+#include "xla/python/ifrt/attribute_map.h"
 
 // IFRT - PJRT
 #include "xla/python/pjrt_ifrt/pjrt_array.h"
@@ -1567,11 +1568,10 @@ ifrt_proxy_grpc_server_create_from_ifrt_client_factory_tpu(
   return MyValueOrThrow(
              xla::ifrt::proxy::GrpcServer::CreateFromIfrtClientFactory(
                  address,
-                 []() -> absl::StatusOr<std::shared_ptr<xla::ifrt::Client>> {
+                 [](xla::ifrt::AttributeMap initialization_data) -> absl::StatusOr<std::shared_ptr<xla::ifrt::Client>> {
                    auto pjrt_client =
                        std::shared_ptr<xla::PjRtClient>(GetCApiClient("TPU"));
-                   return std::shared_ptr<xla::ifrt::Client>(
-                       xla::ifrt::PjRtClient::Create(pjrt_client).release());
+                   return xla::ifrt::PjRtClient::Create(std::move(pjrt_client));
                  }))
       .release();
 }
diff --git a/deps/ReactantExtra/WORKSPACE b/deps/ReactantExtra/WORKSPACE
index 2f17416637..353ebacea4 100644
--- a/deps/ReactantExtra/WORKSPACE
+++ b/deps/ReactantExtra/WORKSPACE
@@ -129,7 +129,9 @@ http_archive(
     patches = ["@enzyme_ad//:patches/jax.patch"],
 )
 
-load("@jax//third_party/xla:workspace.bzl", "XLA_COMMIT", "XLA_SHA256")
+# load("@jax//third_party/xla:workspace.bzl", "XLA_COMMIT", "XLA_SHA256")
+XLA_COMMIT = "b34ff92fce1c51bf895a80fb1a3672b1301eb7aa"
+XLA_SHA256 = ""
 
 http_archive(
     name = "xla",
diff --git a/src/Compiler.jl b/src/Compiler.jl
index 16aba88de6..c26b625029 100644
--- a/src/Compiler.jl
+++ b/src/Compiler.jl
@@ -645,7 +645,8 @@ function compile_mlir!(
     end
 
     raise = if Raise[]
-        "llvm-to-memref-access,canonicalize,convert-llvm-to-cf,canonicalize,enzyme-lift-cf-to-scf,canonicalize,func.func(canonicalize-loops),canonicalize-scf-for,canonicalize,affine-cfg,canonicalize,func.func(canonicalize-loops),canonicalize,llvm-to-affine-access,canonicalize,delinearize-indexing,canonicalize"
+        # "llvm-to-memref-access" # ,canonicalize,convert-llvm-to-cf,canonicalize,enzyme-lift-cf-to-scf,canonicalize,func.func(canonicalize-loops),canonicalize-scf-for,canonicalize,affine-cfg,canonicalize,func.func(canonicalize-loops),canonicalize,llvm-to-affine-access,canonicalize,delinearize-indexing,canonicalize"
+        "canonicalize"
     else
         "canonicalize"
     end