Skip to content

Commit b52fecc

Browse files
authored
Merge pull request #308 from regro-cf-autotick-bot/0.5.3_h341871
jaxlib v0.5.3
2 parents 998de78 + c81b0f4 commit b52fecc

File tree

6 files changed

+58
-111
lines changed

6 files changed

+58
-111
lines changed

recipe/build.sh

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -102,11 +102,6 @@ fi
102102
# For debugging
103103
# CUSTOM_BAZEL_OPTIONS="${CUSTOM_BAZEL_OPTIONS} --bazel_options=--subcommands"
104104

105-
if [[ "${target_platform}" == "osx-64" ]]; then
106-
# Tensorflow doesn't cope yet with an explicit architecture (darwin_x86_64) on osx-64 yet.
107-
TARGET_CPU=darwin
108-
fi
109-
110105
# Force static linkage with protobuf to avoid definition collisions,
111106
# see https://github.com/conda-forge/jaxlib-feedstock/issues/89
112107
#
@@ -118,6 +113,8 @@ if [[ "${target_platform}" == "osx-arm64" || "${target_platform}" != "${build_pl
118113
else
119114
EXTRA="${CUDA_ARGS:-}"
120115
fi
116+
# Never use the Appe toolchain
117+
sed -i '/local_config_apple/d' .bazelrc
121118
if [[ "${target_platform}" == linux-* ]]; then
122119
EXTRA="${EXTRA} --use_clang false --gcc_path $CC"
123120

recipe/meta.yaml

Lines changed: 5 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
1-
{% set version = "0.5.2" %}
2-
{% set build = 2 %}
1+
{% set version = "0.5.3" %}
2+
{% set build = 0 %}
33

44
{% if cuda_compiler_version != "None" %}
55
{% set build = build + 200 %}
@@ -13,14 +13,10 @@ package:
1313
source:
1414
# only pull sources after upstream PyPI release...
1515
url: https://github.com/jax-ml/jax/archive/jax-v{{ version }}.tar.gz
16-
sha256: 8e9de1e012dd65fc4a9eec8af4aa2bf6782767130a5d8e1c1e342b7d658280fe
16+
sha256: 1094581a30ec069965f4e3e67d60262570cc3dd016adc62073bc24347b14270c
1717
patches:
1818
- patches/0001-Allow-for-custom-CUDA-build.patch
1919
- patches/0002-Consolidated-build-fixes-for-XLA.patch
20-
# backport https://github.com/jax-ml/jax/pull/26691
21-
- patches/0003-fix-member-access-to-packed-CUDA-struct.patch
22-
# backport https://github.com/jax-ml/jax/pull/26690
23-
- patches/0004-fix-getting-gcc-major-version.patch
2420

2521
build:
2622
number: {{ build }}
@@ -45,8 +41,8 @@ requirements:
4541
- rsync # [(cuda_compiler_version or "").startswith("12")]
4642
# Keep bazel listed twice here to help the migrators track dependencies
4743
- bazel
48-
- bazel 6.*
49-
- bazel-toolchain >=0.1.10
44+
- bazel 7.*
45+
- bazel-toolchain >=0.3.0
5046
# need protoc
5147
- libprotobuf
5248
# needs protoc-gen-grpc

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

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,17 @@
1-
From c97dee8337618e173ed5d617eb67deabdda026c5 Mon Sep 17 00:00:00 2001
1+
From bbdf2dde4f3a32a2265efcb9f7382dd985e7cadc Mon Sep 17 00:00:00 2001
22
From: "Uwe L. Korn" <[email protected]>
33
Date: Sun, 8 Oct 2023 19:34:34 +0200
4-
Subject: [PATCH 1/4] Allow for custom CUDA build
4+
Subject: [PATCH 1/2] Allow for custom CUDA build
55

66
---
77
build/build.py | 14 ++++++++++----
88
1 file changed, 10 insertions(+), 4 deletions(-)
99

1010
diff --git a/build/build.py b/build/build.py
11-
index 0df7d646f..6b9690ca1 100755
11+
index d38b911bb..932587ba8 100755
1212
--- a/build/build.py
1313
+++ b/build/build.py
14-
@@ -549,6 +549,13 @@ async def main():
14+
@@ -573,6 +573,13 @@ async def main():
1515

1616
if args.cuda_version:
1717
logging.debug("Hermetic CUDA version: %s", args.cuda_version)
@@ -25,9 +25,9 @@ index 0df7d646f..6b9690ca1 100755
2525
wheel_build_command_base.append(
2626
f"--repo_env=HERMETIC_CUDA_VERSION={args.cuda_version}"
2727
)
28-
@@ -656,10 +663,9 @@ async def main():
28+
@@ -687,10 +694,9 @@ async def main():
2929

30-
wheel_build_command.append(f"--jaxlib_git_hash={git_hash}")
30+
wheel_build_command.append(f"--jaxlib_git_hash={git_hash}")
3131

3232
- result = await executor.run(wheel_build_command.get_command_as_string(), args.dry_run, args.detailed_timestamped_log)
3333
- # Exit with error if any wheel build fails.

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

Lines changed: 45 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
1-
From 1dd800f6f64d0948feccef00840d74b5c61094e0 Mon Sep 17 00:00:00 2001
1+
From a3b58e63437e9c2432e5704dbb4c97040fd49a81 Mon Sep 17 00:00:00 2001
22
From: "Uwe L. Korn" <[email protected]>
33
Date: Thu, 14 Dec 2023 17:06:15 +0100
4-
Subject: [PATCH 2/4] Consolidated build fixes for XLA
4+
Subject: [PATCH 2/2] Consolidated build fixes for XLA
55

66
jax vendors xla, but only populates the sources through bazel, so we cannot
77
patch as usual through conda, but rather need to teach the bazel build file
@@ -29,11 +29,11 @@ Co-Authored-By: H. Vetinari <[email protected]>
2929

3030
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
3131
new file mode 100644
32-
index 000000000..8554f5126
32+
index 000000000..a320c4b45
3333
--- /dev/null
3434
+++ b/third_party/xla/0001-Support-third-party-build-of-boringssl.patch
3535
@@ -0,0 +1,51 @@
36-
+From c6bf0f14e159dcbdc99c725823d23e1d26b73a04 Mon Sep 17 00:00:00 2001
36+
+From 2b462dc604d15100ece5e942f29a8a2194459c62 Mon Sep 17 00:00:00 2001
3737
+From: "Uwe L. Korn" <[email protected]>
3838
+Date: Thu, 14 Dec 2023 15:04:51 +0100
3939
+Subject: [PATCH 1/5] Support third-party build of boringssl
@@ -72,7 +72,7 @@ index 000000000..8554f5126
7272
++ ],
7373
++)
7474
+diff --git a/workspace2.bzl b/workspace2.bzl
75-
+index cded32c8ca..5fb10ca45f 100644
75+
+index 8a0336e338..870a613377 100644
7676
+--- a/workspace2.bzl
7777
++++ b/workspace2.bzl
7878
+@@ -108,7 +108,7 @@ def _tf_repositories():
@@ -86,11 +86,11 @@ index 000000000..8554f5126
8686
+
8787
diff --git a/third_party/xla/0002-Fix-abseil-headers.patch b/third_party/xla/0002-Fix-abseil-headers.patch
8888
new file mode 100644
89-
index 000000000..7b2084d01
89+
index 000000000..14c782433
9090
--- /dev/null
9191
+++ b/third_party/xla/0002-Fix-abseil-headers.patch
9292
@@ -0,0 +1,73 @@
93-
+From 1620a0eb9135e380b59583539080337853c130aa Mon Sep 17 00:00:00 2001
93+
+From a6d8f25aaa42bf52d40259be4300befc2abd7b20 Mon Sep 17 00:00:00 2001
9494
+From: "Uwe L. Korn" <[email protected]>
9595
+Date: Thu, 23 May 2024 15:45:52 +0200
9696
+Subject: [PATCH 2/5] Fix abseil headers
@@ -103,7 +103,7 @@ index 000000000..7b2084d01
103103
+ 4 files changed, 10 insertions(+)
104104
+
105105
+diff --git a/xla/python/ifrt_proxy/common/BUILD b/xla/python/ifrt_proxy/common/BUILD
106-
+index 5da3b67ecb..c684b54d52 100644
106+
+index ce0c1b6910..9ee4427cc8 100644
107107
+--- a/xla/python/ifrt_proxy/common/BUILD
108108
++++ b/xla/python/ifrt_proxy/common/BUILD
109109
+@@ -51,6 +51,9 @@ cc_library(
@@ -129,10 +129,10 @@ index 000000000..7b2084d01
129129
+ "@tsl//tsl/platform:protobuf",
130130
+ "@tsl//tsl/platform:stringpiece",
131131
+diff --git a/xla/tsl/platform/default/BUILD b/xla/tsl/platform/default/BUILD
132-
+index 829cef6d2a..ae97fe2b92 100644
132+
+index 56e7914198..d72ddf5c7a 100644
133133
+--- a/xla/tsl/platform/default/BUILD
134134
++++ b/xla/tsl/platform/default/BUILD
135-
+@@ -237,6 +237,8 @@ cc_library(
135+
+@@ -238,6 +238,8 @@ cc_library(
136136
+ deps = [
137137
+ "//xla/tsl/platform:logging",
138138
+ "@com_google_absl//absl/log:check",
@@ -142,10 +142,10 @@ index 000000000..7b2084d01
142142
+ )
143143
+
144144
+diff --git a/xla/tsl/profiler/rpc/client/BUILD b/xla/tsl/profiler/rpc/client/BUILD
145-
+index c2554c43cd..5430b49357 100644
145+
+index aed51b8521..c11a6096b1 100644
146146
+--- a/xla/tsl/profiler/rpc/client/BUILD
147147
++++ b/xla/tsl/profiler/rpc/client/BUILD
148-
+@@ -99,6 +99,8 @@ cc_library(
148+
+@@ -102,6 +102,8 @@ cc_library(
149149
+ deps = [
150150
+ ":profiler_client_impl",
151151
+ "//xla/tsl/platform:status",
@@ -154,7 +154,7 @@ index 000000000..7b2084d01
154154
+ "@com_google_absl//absl/strings",
155155
+ "@com_google_absl//absl/time",
156156
+ "@tsl//tsl/profiler/protobuf:profiler_analysis_cc_grpc_proto",
157-
+@@ -126,6 +128,8 @@ cc_library(
157+
+@@ -130,6 +132,8 @@ cc_library(
158158
+ "//xla/tsl/platform:status",
159159
+ "//xla/tsl/platform:types",
160160
+ "//xla/tsl/protobuf:error_codes_proto_impl_cc",
@@ -165,11 +165,11 @@ index 000000000..7b2084d01
165165
+ "@com_google_absl//absl/time",
166166
diff --git a/third_party/xla/0003-Omit-usage-of-StrFormat.patch b/third_party/xla/0003-Omit-usage-of-StrFormat.patch
167167
new file mode 100644
168-
index 000000000..ac7be8f59
168+
index 000000000..dd1057aac
169169
--- /dev/null
170170
+++ b/third_party/xla/0003-Omit-usage-of-StrFormat.patch
171171
@@ -0,0 +1,25 @@
172-
+From 7240379e230c6095e8e490aa8c00e9dc97c0c4f8 Mon Sep 17 00:00:00 2001
172+
+From 91f1fb8157992964d54ad5197247d36cb304151d Mon Sep 17 00:00:00 2001
173173
+From: "Uwe L. Korn" <[email protected]>
174174
+Date: Thu, 4 Jul 2024 10:36:03 +0200
175175
+Subject: [PATCH 3/5] Omit usage of StrFormat
@@ -196,30 +196,30 @@ index 000000000..ac7be8f59
196196
+ #include <vector>
197197
diff --git a/third_party/xla/0004-Add-missing-bits-absl-systemlib.patch b/third_party/xla/0004-Add-missing-bits-absl-systemlib.patch
198198
new file mode 100644
199-
index 000000000..d1ffe736c
199+
index 000000000..e4380aa60
200200
--- /dev/null
201201
+++ b/third_party/xla/0004-Add-missing-bits-absl-systemlib.patch
202202
@@ -0,0 +1,262 @@
203-
+From fa66d4f7f1597cd20962cae167c0a9f09080954b Mon Sep 17 00:00:00 2001
203+
+From 3249c3354be3e9ab9f22f9ba1bcff3900f3bffe6 Mon Sep 17 00:00:00 2001
204204
+From: "Uwe L. Korn" <[email protected]>
205205
+Date: Thu, 4 Jul 2024 15:58:32 +0200
206206
+Subject: [PATCH 4/5] Add missing bits absl systemlib
207207
+
208208
+---
209-
+ .../third_party/absl/system.absl.base.BUILD | 16 ++++
210-
+ .../third_party/absl/system.absl.crc.BUILD | 70 +++++++++++++++++
211-
+ .../third_party/absl/system.absl.flags.BUILD | 2 +-
212-
+ .../third_party/absl/system.absl.log.BUILD | 78 +++++++++++++++++++
213-
+ .../absl/system.absl.strings.BUILD | 12 +++
214-
+ .../tsl/third_party/absl/workspace.bzl | 2 +
209+
+ third_party/absl/system.absl.base.BUILD | 16 +++++
210+
+ third_party/absl/system.absl.crc.BUILD | 70 +++++++++++++++++++
211+
+ third_party/absl/system.absl.flags.BUILD | 2 +-
212+
+ third_party/absl/system.absl.log.BUILD | 78 ++++++++++++++++++++++
213+
+ third_party/absl/system.absl.strings.BUILD | 12 ++++
214+
+ third_party/absl/workspace.bzl | 2 +
215215
+ 6 files changed, 179 insertions(+), 1 deletion(-)
216-
+ create mode 100644 third_party/tsl/third_party/absl/system.absl.crc.BUILD
217-
+ create mode 100644 third_party/tsl/third_party/absl/system.absl.log.BUILD
216+
+ create mode 100644 third_party/absl/system.absl.crc.BUILD
217+
+ create mode 100644 third_party/absl/system.absl.log.BUILD
218218
+
219-
+diff --git a/third_party/tsl/third_party/absl/system.absl.base.BUILD b/third_party/tsl/third_party/absl/system.absl.base.BUILD
219+
+diff --git a/third_party/absl/system.absl.base.BUILD b/third_party/absl/system.absl.base.BUILD
220220
+index d6bf8748de..b77fe118e3 100644
221-
+--- a/third_party/tsl/third_party/absl/system.absl.base.BUILD
222-
++++ b/third_party/tsl/third_party/absl/system.absl.base.BUILD
221+
+--- a/third_party/absl/system.absl.base.BUILD
222+
++++ b/third_party/absl/system.absl.base.BUILD
223223
+@@ -20,6 +20,14 @@ cc_library(
224224
+ linkopts = ["-labsl_log_severity"],
225225
+ )
@@ -247,11 +247,11 @@ index 000000000..d1ffe736c
247247
++ ":core_headers",
248248
++ ],
249249
++)
250-
+diff --git a/third_party/tsl/third_party/absl/system.absl.crc.BUILD b/third_party/tsl/third_party/absl/system.absl.crc.BUILD
250+
+diff --git a/third_party/absl/system.absl.crc.BUILD b/third_party/absl/system.absl.crc.BUILD
251251
+new file mode 100644
252252
+index 0000000000..487c36ff45
253253
+--- /dev/null
254-
++++ b/third_party/tsl/third_party/absl/system.absl.crc.BUILD
254+
++++ b/third_party/absl/system.absl.crc.BUILD
255255
+@@ -0,0 +1,70 @@
256256
++load("@rules_cc//cc:defs.bzl", "cc_library")
257257
++
@@ -323,10 +323,10 @@ index 000000000..d1ffe736c
323323
++ "//absl/base:config",
324324
++ ],
325325
++)
326-
+diff --git a/third_party/tsl/third_party/absl/system.absl.flags.BUILD b/third_party/tsl/third_party/absl/system.absl.flags.BUILD
326+
+diff --git a/third_party/absl/system.absl.flags.BUILD b/third_party/absl/system.absl.flags.BUILD
327327
+index aff653c7e5..5f95b11dec 100644
328-
+--- a/third_party/tsl/third_party/absl/system.absl.flags.BUILD
329-
++++ b/third_party/tsl/third_party/absl/system.absl.flags.BUILD
328+
+--- a/third_party/absl/system.absl.flags.BUILD
329+
++++ b/third_party/absl/system.absl.flags.BUILD
330330
+@@ -97,7 +97,7 @@ cc_library(
331331
+
332332
+ cc_library(
@@ -336,11 +336,11 @@ index 000000000..d1ffe736c
336336
+ deps = [
337337
+ ":config",
338338
+ ":flag_internal",
339-
+diff --git a/third_party/tsl/third_party/absl/system.absl.log.BUILD b/third_party/tsl/third_party/absl/system.absl.log.BUILD
339+
+diff --git a/third_party/absl/system.absl.log.BUILD b/third_party/absl/system.absl.log.BUILD
340340
+new file mode 100644
341341
+index 0000000000..c23aa51acf
342342
+--- /dev/null
343-
++++ b/third_party/tsl/third_party/absl/system.absl.log.BUILD
343+
++++ b/third_party/absl/system.absl.log.BUILD
344344
+@@ -0,0 +1,78 @@
345345
++load("@rules_cc//cc:defs.bzl", "cc_library")
346346
++
@@ -420,10 +420,10 @@ index 000000000..d1ffe736c
420420
++ "-labsl_log_internal_log_sink_set",
421421
++ ],
422422
++)
423-
+diff --git a/third_party/tsl/third_party/absl/system.absl.strings.BUILD b/third_party/tsl/third_party/absl/system.absl.strings.BUILD
423+
+diff --git a/third_party/absl/system.absl.strings.BUILD b/third_party/absl/system.absl.strings.BUILD
424424
+index fa9a7a84f6..6a898a6d48 100644
425-
+--- a/third_party/tsl/third_party/absl/system.absl.strings.BUILD
426-
++++ b/third_party/tsl/third_party/absl/system.absl.strings.BUILD
425+
+--- a/third_party/absl/system.absl.strings.BUILD
426+
++++ b/third_party/absl/system.absl.strings.BUILD
427427
+@@ -2,6 +2,18 @@ load("@rules_cc//cc:defs.bzl", "cc_library")
428428
+
429429
+ package(default_visibility = ["//visibility:public"])
@@ -443,10 +443,10 @@ index 000000000..d1ffe736c
443443
+ cc_library(
444444
+ name = "strings",
445445
+ linkopts = ["-labsl_strings"],
446-
+diff --git a/third_party/tsl/third_party/absl/workspace.bzl b/third_party/tsl/third_party/absl/workspace.bzl
447-
+index d5973b13b3..0bb3fbbf07 100644
448-
+--- a/third_party/tsl/third_party/absl/workspace.bzl
449-
++++ b/third_party/tsl/third_party/absl/workspace.bzl
446+
+diff --git a/third_party/absl/workspace.bzl b/third_party/absl/workspace.bzl
447+
+index c7f4983081..71bab0e819 100644
448+
+--- a/third_party/absl/workspace.bzl
449+
++++ b/third_party/absl/workspace.bzl
450450
+@@ -14,12 +14,14 @@ def repo():
451451
+ SYS_DIRS = [
452452
+ "algorithm",
@@ -464,11 +464,11 @@ index 000000000..d1ffe736c
464464
+ "numeric",
465465
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
466466
new file mode 100644
467-
index 000000000..271a3ca1e
467+
index 000000000..9aa8fcc80
468468
--- /dev/null
469469
+++ b/third_party/xla/0005-Check-whether-absl-log-is-already-initialized.patch
470470
@@ -0,0 +1,58 @@
471-
+From 9ebd25248eda049702cb781894ca0e097f6843b0 Mon Sep 17 00:00:00 2001
471+
+From fb526c7d4d46b9a7d9323f5b51c429edb0d0adee Mon Sep 17 00:00:00 2001
472472
+From: "Uwe L. Korn" <[email protected]>
473473
+Date: Fri, 22 Nov 2024 10:51:18 +0100
474474
+Subject: [PATCH 5/5] Check whether absl log is already initialized
@@ -527,7 +527,7 @@ index 000000000..271a3ca1e
527527
+
528528
+ } // namespace xla
529529
diff --git a/third_party/xla/workspace.bzl b/third_party/xla/workspace.bzl
530-
index 185ee6a39..a7d083e08 100644
530+
index 73bf2eb38..a00de6486 100644
531531
--- a/third_party/xla/workspace.bzl
532532
+++ b/third_party/xla/workspace.bzl
533533
@@ -30,6 +30,13 @@ def repo():

recipe/patches/0003-fix-member-access-to-packed-CUDA-struct.patch

Lines changed: 0 additions & 24 deletions
This file was deleted.

recipe/patches/0004-fix-getting-gcc-major-version.patch

Lines changed: 0 additions & 22 deletions
This file was deleted.

0 commit comments

Comments
 (0)