Merge pull request #291 from xhochy/0.4.35

jaxlib 0.4.35

recipe/build.sh

@@ -72,6 +72,7 @@ fi
 source gen-bazel-toolchain
 
 cat >> .bazelrc <<EOF
+
 build --crosstool_top=//bazel_toolchain:toolchain
 build --logging=6
 build --verbose_failures

recipe/meta.yaml

@@ -1,5 +1,5 @@
-{% set version = "0.4.34" %}
-{% set build = 2 %}
+{% set version = "0.4.35" %}
+{% set build = 0 %}
 
 {% if cuda_compiler_version != "None" %}
 {% set build = build + 200 %}
@@ -13,7 +13,7 @@ package:
 source:
   # only pull sources after upstream PyPI release...
   url: https://github.com/google/jax/archive/jax-v{{ version }}.tar.gz
-  sha256: d3a75ad667772309ade81350fa70c4a78028a920028800282e46d8383c0ee6bb
+  sha256: 65e086708ae56670676b7b2340ad82b901d8c9993d1241a839c8990bdb8d6212
   patches:
     - patches/0001-Allow-for-custom-CUDA-build.patch
     - patches/0002-Consolidated-build-fixes-for-XLA.patch

recipe/patches/0001-Allow-for-custom-CUDA-build.patch

@@ -1,4 +1,4 @@
-From 868e01e0e6672f0b85d63e5564889131a8052e5f Mon Sep 17 00:00:00 2001
+From 66310dedabfce69aea6814da770a95755ecc3ef2 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/3] Allow for custom CUDA build
@@ -8,7 +8,7 @@ Subject: [PATCH 1/3] Allow for custom CUDA build
  1 file changed, 5 insertions(+)
 
 diff --git a/build/build.py b/build/build.py
-index de0d5a9..8b7a21f 100755
+index 44343eb..0a7a003 100755
 --- a/build/build.py
 +++ b/build/build.py
 @@ -292,6 +292,11 @@ def write_bazelrc(*, remote_build,

recipe/patches/0002-Consolidated-build-fixes-for-XLA.patch

@@ -1,4 +1,4 @@
-From e1ce0773ec91aeb5cad98b7ba7332e287151eb2a Mon Sep 17 00:00:00 2001
+From 0bfc707ab06b3f1cdcb451fb52a4759895938f54 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/3] Consolidated build fixes for XLA
@@ -16,27 +16,29 @@ Co-Authored-By: H. Vetinari <[email protected]>
 ---
  ...pport-third-party-build-of-boringssl.patch |  51 ++++
  third_party/xla/0002-Fix-abseil-headers.patch |  73 ++++++
- .../xla/0003-Omit-usage-of-StrFormat.patch    |  24 ++
+ .../xla/0003-Omit-usage-of-StrFormat.patch    |  25 ++
  ...0004-Add-missing-bits-absl-systemlib.patch | 231 ++++++++++++++++++
  ...ther-absl-log-is-already-initialized.patch |  58 +++++
- third_party/xla/workspace.bzl                 |   7 +
- 6 files changed, 444 insertions(+)
+ .../xla/0006-Add-conda-cuda-path.patch        |  31 +++
+ third_party/xla/workspace.bzl                 |   8 +
+ 7 files changed, 477 insertions(+)
  create mode 100644 third_party/xla/0001-Support-third-party-build-of-boringssl.patch
  create mode 100644 third_party/xla/0002-Fix-abseil-headers.patch
  create mode 100644 third_party/xla/0003-Omit-usage-of-StrFormat.patch
  create mode 100644 third_party/xla/0004-Add-missing-bits-absl-systemlib.patch
  create mode 100644 third_party/xla/0005-Check-whether-absl-log-is-already-initialized.patch
+ create mode 100644 third_party/xla/0006-Add-conda-cuda-path.patch
 
 diff --git a/third_party/xla/0001-Support-third-party-build-of-boringssl.patch b/third_party/xla/0001-Support-third-party-build-of-boringssl.patch
 new file mode 100644
-index 0000000..514ebe1
+index 0000000..dd9de48
 --- /dev/null
 +++ b/third_party/xla/0001-Support-third-party-build-of-boringssl.patch
 @@ -0,0 +1,51 @@
-+From 12f526179896176c1726d5fa9eab2f0332e8ee21 Mon Sep 17 00:00:00 2001
++From a9bb7b6fc4e14025c68596a609d64e4da8960996 Mon Sep 17 00:00:00 2001
 +From: "Uwe L. Korn" <[email protected]>
 +Date: Thu, 14 Dec 2023 15:04:51 +0100
-+Subject: [PATCH 1/5] Support third-party build of boringssl
++Subject: [PATCH 1/6] Support third-party build of boringssl
 +
 +---
 + third_party/boringssl.BUILD | 21 +++++++++++++++++++++
@@ -86,14 +88,14 @@ index 0000000..514ebe1
 + 
 diff --git a/third_party/xla/0002-Fix-abseil-headers.patch b/third_party/xla/0002-Fix-abseil-headers.patch
 new file mode 100644
-index 0000000..e47e8a6
+index 0000000..d3aafc8
 --- /dev/null
 +++ b/third_party/xla/0002-Fix-abseil-headers.patch
 @@ -0,0 +1,73 @@
-+From 87e126882738dd5440430f03ff87fdb4ca31516b Mon Sep 17 00:00:00 2001
++From 3f962611a6b5bbbb9c7d1b5f72d0b6e3feae09cb Mon Sep 17 00:00:00 2001
 +From: "Uwe L. Korn" <[email protected]>
 +Date: Thu, 23 May 2024 15:45:52 +0200
-+Subject: [PATCH 2/5] Fix abseil headers
++Subject: [PATCH 2/6] Fix abseil headers
 +
 +---
 + third_party/tsl/tsl/platform/default/BUILD | 2 ++
@@ -103,7 +105,7 @@ index 0000000..e47e8a6
 + 4 files changed, 10 insertions(+)
 +
 +diff --git a/third_party/tsl/tsl/platform/default/BUILD b/third_party/tsl/tsl/platform/default/BUILD
-+index 8cbb72cfc9..c99667c050 100644
++index 7f18f8108f..94cdfe8540 100644
 +--- a/third_party/tsl/tsl/platform/default/BUILD
 ++++ b/third_party/tsl/tsl/platform/default/BUILD
 +@@ -225,6 +225,8 @@ cc_library(
@@ -142,7 +144,7 @@ index 0000000..e47e8a6
 +         "@tsl//tsl/platform:protobuf",
 +         "@tsl//tsl/platform:status",
 +diff --git a/xla/tsl/profiler/rpc/client/BUILD b/xla/tsl/profiler/rpc/client/BUILD
-+index e135bd82dc..81b79f4600 100644
++index ac61088fd3..2861f77d9a 100644
 +--- a/xla/tsl/profiler/rpc/client/BUILD
 ++++ b/xla/tsl/profiler/rpc/client/BUILD
 +@@ -98,6 +98,8 @@ cc_library(
@@ -154,38 +156,39 @@ index 0000000..e47e8a6
 +         "@com_google_absl//absl/strings",
 +         "@com_google_absl//absl/time",
 +         "@tsl//tsl/platform:status",
-+@@ -121,6 +123,8 @@ cc_library(
-+         "//tensorflow/python/profiler/internal:__pkg__",
++@@ -122,6 +124,8 @@ cc_library(
 +     ]),
 +     deps = [
++         "//xla/tsl/protobuf:error_codes_proto_impl_cc",
 ++        "@com_google_absl//absl/status",
 ++        "@com_google_absl//absl/status:statusor",
 +         "@com_google_absl//absl/memory",
 +         "@com_google_absl//absl/strings",
 +         "@com_google_absl//absl/time",
 diff --git a/third_party/xla/0003-Omit-usage-of-StrFormat.patch b/third_party/xla/0003-Omit-usage-of-StrFormat.patch
 new file mode 100644
-index 0000000..ec44acf
+index 0000000..01fd4df
 --- /dev/null
 +++ b/third_party/xla/0003-Omit-usage-of-StrFormat.patch
-@@ -0,0 +1,24 @@
-+From 62ff28675279b0d72cfe477a94dc1d133d680630 Mon Sep 17 00:00:00 2001
+@@ -0,0 +1,25 @@
++From ac263373f90084807fe80665538ca5018bc3a951 Mon Sep 17 00:00:00 2001
 +From: "Uwe L. Korn" <[email protected]>
 +Date: Thu, 4 Jul 2024 10:36:03 +0200
-+Subject: [PATCH 3/5] Omit usage of StrFormat
++Subject: [PATCH 3/6] Omit usage of StrFormat
 +
 +---
-+ xla/stream_executor/gpu/gpu_executor.h | 4 ++++
-+ 1 file changed, 4 insertions(+)
++ xla/stream_executor/gpu/gpu_executor.h | 5 +++++
++ 1 file changed, 5 insertions(+)
 +
 +diff --git a/xla/stream_executor/gpu/gpu_executor.h b/xla/stream_executor/gpu/gpu_executor.h
-+index 0fdde6096e..6d0825301f 100644
++index 0d90dadf84..25afa7ac9b 100644
 +--- a/xla/stream_executor/gpu/gpu_executor.h
 ++++ b/xla/stream_executor/gpu/gpu_executor.h
-+@@ -18,6 +18,10 @@ limitations under the License.
++@@ -17,6 +17,11 @@ limitations under the License.
++ #define XLA_STREAM_EXECUTOR_GPU_GPU_EXECUTOR_H_
 + 
 + #include <cstdint>
-+ #include <memory>
+++#include <memory>
 ++#include <optional>
 ++#include <string>
 ++#include <sstream>
@@ -195,14 +198,14 @@ index 0000000..ec44acf
 + #include <vector>
 diff --git a/third_party/xla/0004-Add-missing-bits-absl-systemlib.patch b/third_party/xla/0004-Add-missing-bits-absl-systemlib.patch
 new file mode 100644
-index 0000000..2230d05
+index 0000000..cca52f1
 --- /dev/null
 +++ b/third_party/xla/0004-Add-missing-bits-absl-systemlib.patch
 @@ -0,0 +1,231 @@
-+From 0e97115cbe5c1068834f1aa2e4e4956215edca82 Mon Sep 17 00:00:00 2001
++From 1bd70877037d67945c0834aca417a66480416902 Mon Sep 17 00:00:00 2001
 +From: "Uwe L. Korn" <[email protected]>
 +Date: Thu, 4 Jul 2024 15:58:32 +0200
-+Subject: [PATCH 4/5] Add missing bits absl systemlib
++Subject: [PATCH 4/6] Add missing bits absl systemlib
 +
 +---
 + .../third_party/absl/system.absl.base.BUILD   | 16 +++++
@@ -432,14 +435,14 @@ index 0000000..2230d05
 +         "numeric",
 diff --git a/third_party/xla/0005-Check-whether-absl-log-is-already-initialized.patch b/third_party/xla/0005-Check-whether-absl-log-is-already-initialized.patch
 new file mode 100644
-index 0000000..a7f7f0b
+index 0000000..b648cf9
 --- /dev/null
 +++ b/third_party/xla/0005-Check-whether-absl-log-is-already-initialized.patch
 @@ -0,0 +1,58 @@
-+From 4c362c9a88ff5ca06a9a075be12e94f79b9affef Mon Sep 17 00:00:00 2001
++From a4fa7b3d14e3d29ee50a84f4e74e9f76c301ea92 Mon Sep 17 00:00:00 2001
 +From: "Uwe L. Korn" <[email protected]>
 +Date: Fri, 22 Nov 2024 10:51:18 +0100
-+Subject: [PATCH 5/5] Check whether absl log is already initialized
++Subject: [PATCH 5/6] Check whether absl log is already initialized
 +
 +---
 + xla/pjrt/c/pjrt_c_api_gpu.cc | 7 ++++++-
@@ -496,10 +499,19 @@ index 0000000..a7f7f0b
 + }  // namespace xla
 diff --git a/third_party/xla/0006-Add-conda-cuda-path.patch b/third_party/xla/0006-Add-conda-cuda-path.patch
 new file mode 100644
-index 0000000..a7f7f0b
+index 0000000..dcd6a1a
 --- /dev/null
 +++ b/third_party/xla/0006-Add-conda-cuda-path.patch
-@@ -0,0 +1,22 @@
+@@ -0,0 +1,31 @@
++From 79444089903f92a228425a000e47a41b4d9dd810 Mon Sep 17 00:00:00 2001
++From: Silvio Traversaro <[email protected]>
++Date: Tue, 10 Dec 2024 09:42:51 +0100
++Subject: [PATCH 6/6] Add conda cuda path
++
++---
++ .../tsl/tsl/platform/default/cuda_root_path.cc        | 11 +++++++++++
++ 1 file changed, 11 insertions(+)
++
 +diff --git a/third_party/tsl/tsl/platform/default/cuda_root_path.cc b/third_party/tsl/tsl/platform/default/cuda_root_path.cc
 +index ca6da0e553..1d8a9450c0 100644
 +--- a/third_party/tsl/tsl/platform/default/cuda_root_path.cc
@@ -523,7 +535,7 @@ index 0000000..a7f7f0b
 + #endif  // defined(PLATFORM_POSIX) && !defined(__APPLE__)
 + 
 diff --git a/third_party/xla/workspace.bzl b/third_party/xla/workspace.bzl
-index 37861d9..45fbcde 100644
+index 3dd8cbd..1ec0ff4 100644
 --- a/third_party/xla/workspace.bzl
 +++ b/third_party/xla/workspace.bzl
 @@ -30,6 +30,14 @@ def repo():

recipe/patches/0003-Add-missing-typename.patch

@@ -1,4 +1,4 @@
-From fd329cb4cd76ed6221f9ef4f9edddd13fb77d2cf Mon Sep 17 00:00:00 2001
+From 317afc9ff9e2cb021bc81a278710a6a3f147c536 Mon Sep 17 00:00:00 2001
 From: "Uwe L. Korn" <[email protected]>
 Date: Wed, 9 Oct 2024 09:39:46 +0000
 Subject: [PATCH 3/3] Add missing typename
@@ -8,10 +8,10 @@ Subject: [PATCH 3/3] Add missing typename
  1 file changed, 1 insertion(+), 1 deletion(-)
 
 diff --git a/jaxlib/gpu/solver_kernels_ffi.cc b/jaxlib/gpu/solver_kernels_ffi.cc
-index 32cd975..c90dfab 100644
+index 7852da4..b5742b5 100644
 --- a/jaxlib/gpu/solver_kernels_ffi.cc
 +++ b/jaxlib/gpu/solver_kernels_ffi.cc
-@@ -815,7 +815,7 @@ ffi::Error GesvdjImpl(int64_t batch, int64_t rows, int64_t cols,
+@@ -915,7 +915,7 @@ ffi::Error GesvdjImpl(int64_t batch, int64_t rows, int64_t cols,
 
    auto a_data = static_cast<T*>(a.untyped_data());
    auto out_data = static_cast<T*>(out->untyped_data());