conda-forge · xhochy · Jun 26, 2024 · Jun 25, 2024
@@ -1,4 +1,4 @@
-{% set version = "0.4.24" %}
+{% set version = "0.4.25" %}
 {% set number = 0 %}
 
 {% if cuda_compiler_version != "None" %}
@@ -13,12 +13,11 @@ package:
 source:
   # only pull sources after upstream PyPI release...
   url: https://github.com/google/jax/archive/jaxlib-v{{ version }}.tar.gz
-  sha256: c4e6963c2c36f634a9a1765e476a1ed4e6c4a7954465ebf72e29f344c28ddc28
+  sha256: fc1197c401924942eb14185a61688d0c476e3e81ff71f9dc95e620b57c06eec8
   patches:
     - patches/0001-Allow-for-custom-CUDA-build.patch
     - patches/0002-xla-Support-third-party-build-of-boringssl.patch
     - patches/0003-xla-Fix-abseil-headers.patch
-    - patches/0004-Upstream-xla-Patch.patch
 
 build:
   number: {{ number }}
@@ -81,7 +80,6 @@ requirements:
     - zlib
   run:
     - python
-    - {{ pin_compatible('numpy') }}
     - scipy >=1.9
     - ml_dtypes >=0.2.0
     - __cuda  # [cuda_compiler_version != "None"]

@@ -1,7 +1,7 @@
-From 7e9afda0f5ad8a0767f0500eb9933b27dde0414e Mon Sep 17 00:00:00 2001
+From 39bed7f08f495d3a544a11c951869fcfd1616003 Mon Sep 17 00:00:00 2001
 From: "Uwe L. Korn" <[email protected]>
 Date: Sun, 8 Oct 2023 19:34:34 +0200
-Subject: [PATCH 1/4] Allow for custom CUDA build
+Subject: [PATCH 1/3] Allow for custom CUDA build
 
 ---
  build/build.py | 5 +++++

@@ -1,7 +1,7 @@
-From 691dd89b6d14b7c71302a46f54b1564c486ea7ba Mon Sep 17 00:00:00 2001
+From a7ae8b8397a8a3f0f1206c08f3461aff9b686670 Mon Sep 17 00:00:00 2001
 From: "Uwe L. Korn" <[email protected]>
 Date: Thu, 14 Dec 2023 17:06:15 +0100
-Subject: [PATCH 2/4] xla: Support third-party build of boringssl
+Subject: [PATCH 2/3] xla: Support third-party build of boringssl
 
 ---
  ...pport-third-party-build-of-boringssl.patch | 54 +++++++++++++++++++
@@ -70,7 +70,7 @@ index 0000000..1f11444
 +2.43.0
 +
 diff --git a/third_party/xla/workspace.bzl b/third_party/xla/workspace.bzl
-index 0b09fa9..b81e13a 100644
+index ebc8d98..ebb0137 100644
 --- a/third_party/xla/workspace.bzl
 +++ b/third_party/xla/workspace.bzl
 @@ -29,6 +29,9 @@ def repo():

@@ -1,44 +1,59 @@
-From 730f54ecf5bd630de69368e0a1108737289ecb38 Mon Sep 17 00:00:00 2001
+From 67075c53bebacbeeda10271e2d5bcb31d4050964 Mon Sep 17 00:00:00 2001
 From: "Uwe L. Korn" <[email protected]>
 Date: Sun, 11 Feb 2024 20:18:13 +0100
-Subject: [PATCH 3/4] xla: Fix abseil headers
+Subject: [PATCH 3/3] xla: Fix abseil headers
 
 ---
- third_party/xla/0001-Fix-abseil-headers.patch | 47 +++++++++++++++++++
+ third_party/xla/0001-Fix-abseil-headers.patch | 73 +++++++++++++++++++
  third_party/xla/workspace.bzl                 |  1 +
- 2 files changed, 48 insertions(+)
+ 2 files changed, 74 insertions(+)
  create mode 100644 third_party/xla/0001-Fix-abseil-headers.patch
 
 diff --git a/third_party/xla/0001-Fix-abseil-headers.patch b/third_party/xla/0001-Fix-abseil-headers.patch
 new file mode 100644
-index 0000000..41f411e
+index 0000000..ad2cbaa
 --- /dev/null
 +++ b/third_party/xla/0001-Fix-abseil-headers.patch
-@@ -0,0 +1,47 @@
-+From 97ff4aad5cbb51103ff2bd0880af454b324ff6af Mon Sep 17 00:00:00 2001
-+From: "Uwe L. Korn" <uwelk@xhochy.com>
-+Date: Sun, 11 Feb 2024 19:35:02 +0100
+@@ -0,0 +1,73 @@
++From 138a28794618469fbcab7634c5febca42e10fef6 Mon Sep 17 00:00:00 2001
++From: "Uwe L. Korn" <uwe.korn@quantco.com>
++Date: Thu, 23 May 2024 15:45:52 +0200
 +Subject: [PATCH] Fix abseil headers
 +
 +---
 + third_party/tsl/tsl/distributed_runtime/rpc/BUILD | 1 +
++ third_party/tsl/tsl/platform/default/BUILD        | 2 ++
 + third_party/tsl/tsl/profiler/rpc/client/BUILD     | 4 ++++
-+ 2 files changed, 5 insertions(+)
++ xla/python/ifrt_proxy/common/BUILD                | 3 +++
++ 4 files changed, 10 insertions(+)
 +
 +diff --git a/third_party/tsl/tsl/distributed_runtime/rpc/BUILD b/third_party/tsl/tsl/distributed_runtime/rpc/BUILD
-+index 9b52ee1cb..875b575b4 100644
++index ae18059b63..8b2f356a08 100644
 +--- a/third_party/tsl/tsl/distributed_runtime/rpc/BUILD
 ++++ b/third_party/tsl/tsl/distributed_runtime/rpc/BUILD
-+@@ -43,6 +43,7 @@ cc_library(
++@@ -42,6 +42,7 @@ cc_library(
 +         "//tsl/platform:stringprintf",
 +         "//tsl/protobuf:distributed_runtime_payloads_proto_cc",
 +         "@com_google_absl//absl/status",
 ++        "@com_google_absl//absl/status:statusor",
 +         "@com_google_absl//absl/strings:cord",
 +     ] + tsl_grpc_cc_dependencies(),
 + )
++diff --git a/third_party/tsl/tsl/platform/default/BUILD b/third_party/tsl/tsl/platform/default/BUILD
++index 6b0e9a8c8c..d3c406fb7b 100644
++--- a/third_party/tsl/tsl/platform/default/BUILD
+++++ b/third_party/tsl/tsl/platform/default/BUILD
++@@ -221,6 +221,8 @@ cc_library(
++     deps = [
++         "//tsl/platform:logging",
++         "@com_google_absl//absl/log:check",
+++        "@com_google_absl//absl/status",
+++        "@com_google_absl//absl/status:statusor",
++     ] + tsl_grpc_cc_dependencies(),
++ )
++ 
 +diff --git a/third_party/tsl/tsl/profiler/rpc/client/BUILD b/third_party/tsl/tsl/profiler/rpc/client/BUILD
-+index fb3a39cd1..2b93965e8 100644
++index 1dd7d8a227..ab28e814e5 100644
 +--- a/third_party/tsl/tsl/profiler/rpc/client/BUILD
 ++++ b/third_party/tsl/tsl/profiler/rpc/client/BUILD
 +@@ -100,6 +100,8 @@ cc_library(
@@ -50,20 +65,31 @@ index 0000000..41f411e
 +         "@com_google_absl//absl/strings",
 +         "@com_google_absl//absl/time",
 +     ],
-+@@ -128,6 +130,8 @@ cc_library(
++@@ -126,6 +128,8 @@ cc_library(
++         "//tsl/platform:types",
++         "//tsl/profiler/protobuf:profiler_analysis_cc_grpc_proto",
 +         "//tsl/profiler/protobuf:profiler_service_cc_grpc_proto",
+++        "@com_google_absl//absl/status",
+++        "@com_google_absl//absl/status:statusor",
 +         "//tsl/protobuf:error_codes_proto_impl_cc",
 +         "@com_google_absl//absl/memory",
++         "@com_google_absl//absl/strings",
++diff --git a/xla/python/ifrt_proxy/common/BUILD b/xla/python/ifrt_proxy/common/BUILD
++index 8e750cb525..aede7facd0 100644
++--- a/xla/python/ifrt_proxy/common/BUILD
+++++ b/xla/python/ifrt_proxy/common/BUILD
++@@ -51,6 +51,9 @@ cc_library(
++         "@com_github_grpc_grpc//:grpc++",
++         "@com_google_absl//absl/log",
++         "@com_google_absl//absl/log:check",
 ++        "@com_google_absl//absl/status",
 ++        "@com_google_absl//absl/status:statusor",
-+         "@com_google_absl//absl/strings",
-+         "@com_google_absl//absl/time",
-+     ] + tsl_grpc_cc_dependencies(),
-+-- 
-+2.43.0
-+
+++        "@com_google_absl//absl/strings",
++         "@tsl//tsl/platform",
++     ],
++     alwayslink = True,
 diff --git a/third_party/xla/workspace.bzl b/third_party/xla/workspace.bzl
-index b81e13a..9671b9b 100644
+index ebb0137..1a55893 100644
 --- a/third_party/xla/workspace.bzl
 +++ b/third_party/xla/workspace.bzl
 @@ -31,6 +31,7 @@ def repo():