[deps] Update to new Enzyme-JAX (which updates to new JAX and Co.) (#961)

* [deps] Update to new Enzyme-JAX (which updates to new JAX and Co.)

* Fix build

---------

Co-authored-by: Paul Berg <[email protected]>

Author: giordano

deps/ReactantExtra/API.cpp

@@ -1715,8 +1715,10 @@ extern "C" bool ifrt_DeviceIsAddressable(ifrt::Device *device) {
   return device->IsAddressable();
 }
 
-tsl::RCReference<ifrt::DeviceList> ifrt_CreateDeviceListFromDevices(
-    ifrt::Client *client, ifrt::Device **device_list, int32_t num_devices) {
+static xla::ifrt::RCReferenceWrapper<ifrt::DeviceList>
+ifrt_CreateDeviceListFromDevices(ifrt::Client *client,
+                                 ifrt::Device **device_list,
+                                 int32_t num_devices) {
   absl::Span<ifrt::Device *const> devices(device_list, num_devices);
   return client->MakeDeviceList(devices);
 }

deps/ReactantExtra/WORKSPACE

@@ -9,7 +9,7 @@ http_archive(
     urls = ["https://github.com/wsmoses/nsync/archive/{commit}.tar.gz".format(commit = NSYNC_COMMIT)],
 )
 
-ENZYMEXLA_COMMIT = "7d5b22fc4d9e4d9fb28ba422cb3b6bb510ea0cf9"
+ENZYMEXLA_COMMIT = "6116ac04785ad6c0e160b3614fe6cadc19d479c6"
 ENZYMEXLA_SHA256 = ""
 
 http_archive(
@@ -224,9 +224,6 @@ xla_workspace1()
 load("@xla//:workspace0.bzl", "xla_workspace0")
 xla_workspace0()
 
-load("@jax//third_party/flatbuffers:workspace.bzl", flatbuffers = "repo")
-flatbuffers()
-
 load("@jax//jaxlib:jax_python_wheel.bzl", "jax_python_wheel_repository")
 jax_python_wheel_repository(
     name = "jax_wheel",
@@ -235,15 +232,18 @@ jax_python_wheel_repository(
 )
 
 load(
-    "@tsl//third_party/py:python_wheel.bzl",
+    "@xla//third_party/py:python_wheel.bzl",
     "python_wheel_version_suffix_repository",
 )
 python_wheel_version_suffix_repository(
     name = "jax_wheel_version_suffix",
 )
 
+load("@jax//third_party/flatbuffers:workspace.bzl", flatbuffers = "repo")
+flatbuffers()
+
 load(
-    "@tsl//third_party/gpus/cuda/hermetic:cuda_json_init_repository.bzl",
+    "@xla//third_party/gpus/cuda/hermetic:cuda_json_init_repository.bzl",
     "cuda_json_init_repository",
 )
 
@@ -255,7 +255,7 @@ load(
     "CUDNN_REDISTRIBUTIONS",
 )
 load(
-    "@tsl//third_party/gpus/cuda/hermetic:cuda_redist_init_repositories.bzl",
+    "@xla//third_party/gpus/cuda/hermetic:cuda_redist_init_repositories.bzl",
     "cuda_redist_init_repositories",
     "cudnn_redist_init_repository",
 )
@@ -269,21 +269,21 @@ cudnn_redist_init_repository(
 )
 
 load(
-    "@tsl//third_party/gpus/cuda/hermetic:cuda_configure.bzl",
+    "@xla//third_party/gpus/cuda/hermetic:cuda_configure.bzl",
     "cuda_configure",
 )
 
 cuda_configure(name = "local_config_cuda")
 
 load(
-    "@tsl//third_party/nccl/hermetic:nccl_redist_init_repository.bzl",
+    "@xla//third_party/nccl/hermetic:nccl_redist_init_repository.bzl",
     "nccl_redist_init_repository",
 )
 
 nccl_redist_init_repository()
 
 load(
-    "@tsl//third_party/nccl/hermetic:nccl_configure.bzl",
+    "@xla//third_party/nccl/hermetic:nccl_configure.bzl",
     "nccl_configure",
 )