Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[pull] master from tensorflow:master #169

Open
wants to merge 2,298 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2298 commits
Select commit Hold shift + click to select a range
8802ee7
PR #20794: [gpu][ds-fusion] Add handling for offset module in ds-fusi…
shraiysh Jan 28, 2025
bd65d37
[XLA:GPU] Add triton support test for replica-id, partition-id, colle…
thcmbs Jan 28, 2025
f79054f
[XLA:GPU] Add --merge mode to matmul gen perf table tool.
golechwierowicz Jan 28, 2025
1811319
PR #21805: Fix typo in cli flag help message
Zentrik Jan 28, 2025
554f1df
Triton: Fix missing include header
wsmoses Jan 28, 2025
bda43fd
Fix edge case overflow in `tf.raw_ops.RandomIndexShuffle`.
iindyk Jan 28, 2025
43cd81b
[XLA:CPU] Migrate concatenate kernel generation to ConcatenateKernelE…
WillFroom Jan 28, 2025
5e78985
Set expected entry point name to empty string in model buffer append.
LukeBoyer Jan 28, 2025
7ef8b22
[XLA:GPU] Create a pathway between XLA compile-time and runtime to pa…
vwbaker Jan 28, 2025
4b397a2
Avoid a potential extra string allocation when constructing litert::E…
tensorflower-gardener Jan 28, 2025
b44f57a
[pjrt] Added a PjRtMemorySpace* overload for CreateUninitializedBuffer
superbobry Jan 28, 2025
23cc709
Integrate LLVM at llvm/llvm-project@aa65f93b71de
gribozavr Jan 28, 2025
418251b
[XLA] Correct a few typos in the comments of memory_space_assignment …
tensorflower-gardener Jan 28, 2025
ae58809
Simplify HloRunnerPjRt use of `PjRtClient::BufferFromHostLiteral`.
nvgrw Jan 28, 2025
6f23f51
[xla:gpu] [cleanup] Enclose iterable inputs into their own class
vwbaker Jan 28, 2025
8e51584
Use inspect.getfile instead of inspect.getsourcefile
oprypin Jan 28, 2025
5497a3a
#litert Use `LiteRtGetStatusString()` to make error messages clearer.
qukhan Jan 28, 2025
5c37c69
Poison "quantized int" specializations of std::is_signed, as a first …
tensorflower-gardener Jan 28, 2025
ea84e46
Make `DynamicShape`, `Sharding`, and `ArraySpec` hashable
junwhanahn Jan 28, 2025
2381174
[XLA:GPU] Add operands and flops to matmul table gen.
golechwierowicz Jan 28, 2025
18e5076
Skip TSAN tests as they are also not supported for GPU.
tensorflower-gardener Jan 28, 2025
7ede782
Add support for CUDA 12.8.0 and CUDNN 9.7.0.
tensorflower-gardener Jan 28, 2025
50a5f24
Delete XLA CPU Kokoro build now that the GitHub Actions CPU build is …
ddunl Jan 28, 2025
8fa71b9
Merge pull request #85608 from Jerry-Ge/master
jpienaar Jan 28, 2025
ac6fc5d
Allow backend-specific `test_migrated_to_hlo_runner_pjrt` tags to be …
nvgrw Jan 28, 2025
34426b7
Change default pending shape func to a fixed code location instead of…
cky9301 Jan 28, 2025
19b0a94
[xla:cpu] Add more sort keywords to correctly order XLA:CPU debug opt…
penpornk Jan 28, 2025
833b7a8
Move file translates to tf lite directory for correct ownership.
rocketas Jan 28, 2025
7637fc5
Only load opencl once to mitigate potential un-thread-safety in the i…
tensorflower-gardener Jan 28, 2025
6d844e9
Add new send/recv channel Megascale stat type.
tensorflower-gardener Jan 28, 2025
20fe64e
Print layout info in verifier in some more places
tensorflower-gardener Jan 28, 2025
97875b9
Switch CollectiveMemoryAllocation to use CreateMemoryAllocator infras…
klucke Jan 28, 2025
8a0fbda
Add NCCL dict entry for CUDA 12.8.0.
tensorflower-gardener Jan 28, 2025
8acf753
Replace `kLiteRtAccelator*` with `kLiteRtAccelerator*`.
qukhan Jan 28, 2025
0a0058a
Remove broken and unused support for allocating MemoryType::kHost mem…
klucke Jan 28, 2025
074130b
Silence msan warning by switching transferring socket addresses to us…
pschuh Jan 28, 2025
527d46a
Unify metadata storage with unified buffer management.
LukeBoyer Jan 28, 2025
707a138
[MHLO] Handle dynamic dimensions in HLO<->MHLO
GleasonK Jan 28, 2025
e1bc56c
Add support for GCS File System in Pywrap Profiler.
cliveverghese Jan 28, 2025
7703ef0
Replace "external_buffer" nomenclature with "op_asset" and leverage t…
LukeBoyer Jan 28, 2025
07dff74
#litert Remove unused include in `litert_expected.h`
qukhan Jan 28, 2025
d23a207
Support CreateBuffersForAsyncHostToDevice with shape_spec + layout in…
pschuh Jan 28, 2025
11b46ff
Reverts a47a28e840cf97148669ba3483cd72e87f0efa5b
seherellis Jan 28, 2025
8908646
Encapsulate edge splitting validation attribute in SourceTargetPairs.
toli-y Jan 28, 2025
42898a8
[XLA:MemoryScheduler] Enable constant deferring by default.
Tongfei-Guo Jan 28, 2025
c0ceab4
[pjrt] Use the `PjRtMemorySpace` version of `BufferFromHostBuffer`
superbobry Jan 28, 2025
96448c7
Remove `use_parameter_layout_on_device`.
nvgrw Jan 28, 2025
3fdfc03
Handle offset tensor buffers with unified buffer management in model …
LukeBoyer Jan 28, 2025
1496729
PR #21375: [ds-fusion] Get While loop analysis with copy fusion
shraiysh Jan 28, 2025
1561ef3
[XLA:GPU] Triton support test - minor fixes
thcmbs Jan 28, 2025
3d3fc2d
Change noisy (and useless) XLA log line to VLOG
rhundt Jan 28, 2025
8b90c77
[XLA:LAYOUT_CONSTRAINTS] Move custom call constraints before backend …
blakehechtman Jan 28, 2025
75a8c0b
Refactor IsXSpaceGrouped to verify that all device and host planes ar…
bmass02 Jan 28, 2025
a5c96a2
Update users of TSL headers and targets to new location in XLA
ddunl Jan 28, 2025
25a4f93
Update HLO Translate to return errors using MLIR diagnostics, update …
GleasonK Jan 28, 2025
7bc57ad
Integrate StableHLO at openxla/stablehlo@c27ba678
ghpvnist Jan 28, 2025
9697a39
[pjrt] Use the `PjRtMemorySpace` version of `BufferFromHostBuffer`
superbobry Jan 28, 2025
88b7655
Implement on_device_shape() and logical_on_device_shape().
matthiaskramm Jan 28, 2025
170789c
add PJRT runtime profiling library
tensorflower-gardener Jan 28, 2025
5ee01b7
Improve MSA readabilty by changing some auto types to their type name.
sparc1998 Jan 28, 2025
5e3d02f
Move TSL's `workspace*.bzl` files to XLA
ddunl Jan 28, 2025
5979eac
Add IFRT wrappers to public API
tensorflower-gardener Jan 28, 2025
e2908ff
Use `HloRunnerPropertyTag::kCpu` in convert_test.
nvgrw Jan 28, 2025
dcb7b83
Don't run thread sans on apply plugin test. In general we can't use s…
LukeBoyer Jan 28, 2025
5647b7f
Support offset tensors in model load. This finishes offset tensor sup…
LukeBoyer Jan 29, 2025
0de2fb6
Allow `GetModelBufWithByteCode` to work with models that contain more…
turbotoribio Jan 29, 2025
d00af8d
Use efficient packed flatbuffer api to handle underlying tfl models.
LukeBoyer Jan 29, 2025
f4a9e97
Set CPU as the default acceleration option for a Compiled Model
tensorflower-gardener Jan 29, 2025
6ba7e4e
Actually copy the extra wheels into the expected directory afterwards.
belitskiy Jan 29, 2025
68c0a8b
Make TensorBuffer owner of the associated Event
tensorflower-gardener Jan 29, 2025
451f3a7
Merge pull request #83343 from cybersupersoap:MapUnstage-abort-fix
tensorflower-gardener Jan 29, 2025
1bdabaf
Always dispatch CPU executables synchronously when they include callb…
dfm Jan 29, 2025
50c4e0a
Migrate convolution_test to always use PjRt for its test backend.
nvgrw Jan 29, 2025
0092ab7
Add async/c/types.h to core/c:headers_filegroup
terryheo Jan 29, 2025
bf2d6e7
Migrate array_elementwise_ops_test to always use PjRt for its test ba…
nvgrw Jan 29, 2025
a26f99e
PR #21822: [XLA:GPU] Add support for SM100a architecture (Blackwell)
sergey-kozub Jan 29, 2025
93f10cd
litert: Refine c/litert_dispatch_delegate.h
terryheo Jan 29, 2025
e333034
Automated Code Change
tensorflower-gardener Jan 29, 2025
477c4c3
[xla:cpu] Extract convolution_lib from a convolution_thunk
ezhulenev Jan 29, 2025
3d64e3d
Automated Code Change
tensorflower-gardener Jan 29, 2025
ce10886
Automated Code Change
tensorflower-gardener Jan 29, 2025
700be37
Automated Code Change
tensorflower-gardener Jan 29, 2025
380d9a6
Automated Code Change
tensorflower-gardener Jan 29, 2025
64078fe
Move CudaComputeCapability version parser into factory function
beckerhe Jan 29, 2025
db3bccd
PR #21960: Updated nanobind commit
vfdev-5 Jan 29, 2025
2bbc047
PR #21813: Support e8m0fnu for NCCL collectives
wenscarl Jan 29, 2025
3bf73c7
Reverts 14967296f7faae46550bf5d3e240194e74ebc552
dimitar-asenov Jan 29, 2025
83ae603
Update GraphDef version to 2122.
tensorflower-gardener Jan 29, 2025
7bbce29
compat: Update forward compatibility horizon to 2025-01-29
tensorflower-gardener Jan 29, 2025
90cfd53
Log fused computation from Triton emitter as well
gflegar Jan 29, 2025
9c4cde6
[pjrt] Use the `PjRtMemorySpace*` version of `BufferFromHostLiteral`
superbobry Jan 29, 2025
e7d367e
[XLA:GPU] Turn LHS on when PGO profile is present.
golechwierowicz Jan 29, 2025
2231262
Add a reproducer for issue with reduce->bitcast->broadcast->slice pat…
gflegar Jan 29, 2025
18b9e08
Fix bug in AlgebraicSimplifier
akuegel Jan 29, 2025
1a992b5
[XLA:GPU] Add sweep over batch sizes.
golechwierowicz Jan 29, 2025
b6cf3fc
Automated Code Change
tensorflower-gardener Jan 29, 2025
d40a985
Disable async dispatch within the body of a host callback.
dfm Jan 29, 2025
c09527b
Automated Code Change
tensorflower-gardener Jan 29, 2025
73a09e0
Automated Code Change
tensorflower-gardener Jan 29, 2025
395000c
Parse XLA_FLAGS environment variable every time.
tensorflower-gardener Jan 29, 2025
7ab967c
[XLA:CPU] Integrating ObjectLoader into JITCompiler
tensorflower-gardener Jan 29, 2025
5764746
Automated Code Change
tensorflower-gardener Jan 29, 2025
ca4d9e0
Automated Code Change
tensorflower-gardener Jan 29, 2025
b566131
Integrate LLVM at llvm/llvm-project@a06c89387621
gribozavr Jan 29, 2025
46f98fe
Reverts 3bf73c73e338b45a5e849808e53c2494508d415a
dimitar-asenov Jan 29, 2025
f6988ee
PR #21800: [XLA:GPU] Add block scaling rewriter pass
sergey-kozub Jan 29, 2025
ad67e8c
Reverts 9dcbe84fd12c1a2c48925dfd5be954fad6d9c864
mooskagh Jan 29, 2025
205a086
Create "internal" visibility for LiteRT-internal targets
turbotoribio Jan 29, 2025
6bafbae
[xla:gpu] [cleanup] Pull out some logic into IterableInput
vwbaker Jan 29, 2025
2db6835
[XLA:GPU] Introduce a separate option which would control falling bac…
mooskagh Jan 29, 2025
40e02fe
cleanup of deprecated test methods
tensorflower-gardener Jan 29, 2025
a693ef7
Extend timeout for running E2E XLA CPU benchmarks
juliagmt-google Jan 29, 2025
e6e0261
[xla:ffi] Add support for passing RunId to FFI handlers
ezhulenev Jan 29, 2025
904f490
Update users of TSL headers and targets to new location in XLA
ddunl Jan 29, 2025
079cf0b
[XLA:GPU] Allow AUTO layout in the multihost runner.
mooskagh Jan 29, 2025
278b3ca
Fix collective memory allocation in cuda executor
beckerhe Jan 29, 2025
fa148fd
Fix line wraps in gpu_latency_hiding_scheduler_test
frgossen Jan 29, 2025
965349d
[XLA:CPU][roll forward] Read thunks from proto when loading executable.
tensorflower-gardener Jan 29, 2025
0ea72e8
[Checkpoint] Add units (microseconds) to duration log message.
mrry Jan 29, 2025
326a8c4
Update XNNPack to the lastest version (and dependencies).
qukhan Jan 29, 2025
bd56220
PR #21901: Ci add rocm6.1 deps for ubuntu 20.04
alekstheod Jan 29, 2025
162ac0d
Update users of TSL headers and targets to new location in XLA
ddunl Jan 29, 2025
97cbb04
litert: Fix LiteRT Runtime C library build
terryheo Jan 29, 2025
da05af8
Update Build files to allow visibility to for internal use_case
turbotoribio Jan 29, 2025
cda8b1b
[xla:cpu] Add benchmarks for transpose + copy + dot
ezhulenev Jan 29, 2025
814a500
PR #21948: [GPU] Upgrade cuDNN frontend to 1.10.0.
sergachev Jan 29, 2025
469dfee
Reverts 576474666afd9fd7efa409d56edd48027db4b169
junjiang-lab Jan 29, 2025
ac804e0
PR #22029: [XLA:GPU] Add support for SM101a and SM120a architectures …
sergey-kozub Jan 29, 2025
76eed36
Fix tsan/asan error
deqiangc Jan 29, 2025
baa09a3
Only apply collective permute decomposer to collective-permutes at th…
frgossen Jan 29, 2025
7df3e60
[xla:cpu] Move convolution shape verification and canonical dims comp…
ezhulenev Jan 29, 2025
5541ad9
Add `--ppc64le-compliance-tag` mandatory argument to manylinux_compli…
tensorflower-gardener Jan 29, 2025
f4088c3
Add an ICYU pragma to silence linters.
rhundt Jan 29, 2025
b97365e
[pjrt] `PjRtStreamExecutorDevice` now has a default memory space
superbobry Jan 29, 2025
0ad4c56
Add support for MemoryType::kHost to CreateStreamExecutor on the rele…
klucke Jan 29, 2025
61ee094
Bump up arm64 Linux and Mac wheel size limits: 240M -> 245M.
belitskiy Jan 29, 2025
74c7c8b
Merge pull request #80498 from tensorflow:TFLite-Code-Fix
tensorflower-gardener Jan 29, 2025
134d7a0
[pjrt] Use the `PjRtMemorySpace` version of `BufferFromHostBuffer`
superbobry Jan 29, 2025
ef816eb
Add debug logging
tensorflower-gardener Jan 29, 2025
d5e7459
Disable creation of issues that don't follow the template
mihaimaruseac Jan 29, 2025
0f26341
Merge pull request #82210 from tensorflow:fix/tflite-num-threads-vali…
tensorflower-gardener Jan 29, 2025
51c6eed
[cpp23] Remove std::aligned_storage<> in tensorflow
tensorflower-gardener Jan 29, 2025
e356b81
Make GPU PJRT use CreateMemoryAllocator for Host memory.
klucke Jan 29, 2025
95b4f9d
Modify the TFLite delegate for XNNPACK to use the new `xnn_define_sta…
tensorflower-gardener Jan 29, 2025
5606d32
Use StreamExecutorAllocator instead of DeviceHostAllocator.
klucke Jan 29, 2025
d4b5c66
Update users of TSL headers and targets to new location in XLA
ddunl Jan 30, 2025
e2abe7a
Add single value get_*, add_* and set_front_end_attribute functions t…
toli-y Jan 30, 2025
2d969b4
[GPU] Fix compilation with NVIDIA driver 570.
sergachev Jan 30, 2025
4b6001d
Integrate StableHLO at openxla/stablehlo@48a1e14e
abhigunj Jan 30, 2025
278d556
litert: Tune link flags of LiteRT Runtime
terryheo Jan 30, 2025
531b3d9
Update users of TSL headers and targets to new location in XLA
ddunl Jan 30, 2025
1a227be
Update users of TSL headers and targets to new location in XLA
ddunl Jan 30, 2025
ddceb96
Add GitHub Actions based XLA ARM64 CPU build
ddunl Jan 30, 2025
e11579a
Add license stanza
mihaimaruseac Jan 30, 2025
2d6cc81
Merged commit includes the following changes:
tensorflower-gardener Jan 30, 2025
55f6ed8
Update users of TSL headers and targets to new location in XLA
ddunl Jan 30, 2025
b4f6acf
Disable cp cycle decomposition for pipeline parallelism implementatio…
frgossen Jan 30, 2025
e33b75b
Update users of TSL headers and targets to new location in XLA
ddunl Jan 30, 2025
0e749da
Return arrays from `ArrayImpl._check_and_rearrange`.
emilyfertig Jan 30, 2025
f0a5244
Merge pull request #86164 from tensorflow:mihaimaruseac-patch-1
tensorflower-gardener Jan 30, 2025
d02727c
Generalize cycle related functions and move to SourceTargetPair class.
toli-y Jan 30, 2025
b7a1cfd
Reverts 8b90c7755e298136842b0a952ace340e5b535d23
tensorflower-gardener Jan 30, 2025
2e4a5f9
Disable pipeline parallelism test, which runs into deadlock
frgossen Jan 30, 2025
b2b92ee
Adding a test for strided slice which triggers a path fixed earlier
tensorflower-gardener Jan 30, 2025
7f8fa80
Update users of TSL headers and targets to new location in XLA
ddunl Jan 30, 2025
75b7680
[MultiHostHloRunner] Introduce GPURunnerProfiler to enable GPU profil…
juliagmt-google Jan 30, 2025
35fbed2
Add some additional args for select_and_scatter_test for certain back…
tensorflower-gardener Jan 30, 2025
ec25b45
#tflite Fix compilation of experimental tests.
qukhan Jan 30, 2025
f179097
Make the TF shape inference update the XlaCallModule's StableHLO modu…
eunjaekim-0 Jan 30, 2025
53927a2
[XLA:CPU] Rename `thunk_serdes_proto` to `thunk_proto_serdes`
tensorflower-gardener Jan 30, 2025
28f8d3c
Wrap HLO strings in collective permute decomposer
frgossen Jan 30, 2025
89f71b6
Don't fail on ./dist directory having already been created.
belitskiy Jan 30, 2025
4f7f0db
[pjrt] Migrated remaining uses of `CreateUninitializedBuffer` and `Cr…
superbobry Jan 30, 2025
b5ccfcb
Fix 03 broken links in index.md
gaikwadrahul8 Jan 30, 2025
508bccc
- Run E2E CPU benchmarks and microbenchmarks on Linux x86 CPU and Lin…
juliagmt-google Jan 30, 2025
7e6783f
[xla:cpu] Add XnnConvolutionThunk
ezhulenev Jan 30, 2025
85f60df
Only decompose collective permute ops that may be pipelined
frgossen Jan 30, 2025
b134803
Add transformation to allow transposed input to BatchMatMul
vamsimanchala Jan 30, 2025
a9bc57b
Fix rare crash in memcpy a2a.
reedwm Jan 30, 2025
60e37ec
Fix `TfrtCpuClient` allocating oversized buffers.
nvgrw Jan 30, 2025
3995d24
Merge pull request #86234 from tensorflow:gaikwadrahul8-patch-1
tensorflower-gardener Jan 30, 2025
81b0a48
Reverts a5c0385cb45c4fbaab7ee2edc0884061be684eb8
kenfranko Jan 30, 2025
f27ff24
[StableHLO] Remove XlaCallModule's MHLO dependency
GleasonK Jan 30, 2025
e05f2be
Call Cleanup in thunks holding other thunks, like SequentialThunk.
reedwm Jan 30, 2025
20b2a63
Add todo to add additional control dependencies
frgossen Jan 30, 2025
0c3b524
Add an option to the TF standard pipeline to enable the StableHLO sha…
eunjaekim-0 Jan 30, 2025
22a3252
Use GenericMemoryAllocation instead of HostMemoryAllocation.
klucke Jan 30, 2025
aa7ce33
Reverts 395000c6a82bdfb51b6ac24c9e9d649556bf292a
hawkinsp Jan 30, 2025
cc7101f
Migrate convert_test to always use PjRt for its test backend.
nvgrw Jan 31, 2025
837ae82
Add documentation for regenerating `warnings.bazelrc`
ddunl Jan 31, 2025
151fc4a
litert: Use BuiltinOpResolver to enable lazy applying Xnnpack delegate
terryheo Jan 31, 2025
5f3663d
Reverts 42898a878b1fd2269a7de1ea22c6cfa1c8914247
tensorflower-gardener Jan 31, 2025
41ea5d3
litert: Do not strip symbol for debug build
terryheo Jan 31, 2025
083c0e1
Add async support for NPU
tensorflower-gardener Jan 31, 2025
a48d834
Remove now unused StreamExecutor::HostMemoryDeallocate method and ove…
klucke Jan 31, 2025
892ba95
Do not fuse instructions inside custom fusions/calls
tensorflower-gardener Jan 31, 2025
0c05718
[XLA:GPU] remove Assert(No)Transform helper methods
tensorflower-gardener Jan 31, 2025
0ffcbd2
[xla:cpu] Remove ACL support from convolution thunk
ezhulenev Jan 31, 2025
e7d2d73
Update TF CUDA 12 Dockerfiles for Python 3.13-nogil support
kanglant Jan 31, 2025
6d5ee9a
[xla:cpu] Don't use deprecated AsyncValueRef::SetError API
ezhulenev Jan 31, 2025
8d6bd5d
Integrate LLVM at llvm/llvm-project@3a1e157454ec
slackito Jan 31, 2025
4af5c5a
Automated Code Change
tensorflower-gardener Jan 31, 2025
8b6d296
Automated Code Change
tensorflower-gardener Jan 31, 2025
4e19439
compat: Update forward compatibility horizon to 2025-01-31
tensorflower-gardener Jan 31, 2025
e3d145e
Update GraphDef version to 2124.
tensorflower-gardener Jan 31, 2025
1ed4981
Rolling integrate Triton up to [515467a9](https://github.com/openai/t…
gflegar Jan 31, 2025
55ef4e2
PR #22101: Define SO_ZEROCOPY, SO_EE_ORIGIN_ZEROCOPY when missing (e.…
pearu Jan 31, 2025
8800e14
Automated Code Change
tensorflower-gardener Jan 31, 2025
f1865d8
Automated Code Change
tensorflower-gardener Jan 31, 2025
67ba087
[XLA:GPU] Prepare add_dim lambda to become a member function.
loislo Jan 31, 2025
58a7342
[XLA/Triton] Don't restrict contracting dimension tiling for predicat…
Moerafaat Jan 31, 2025
50202c4
[XLA:GPU] Convert add_dim lambda to to a private method AddDimToTenso…
loislo Jan 31, 2025
2333637
[XLA:CPU] Enable setting optimization level of testlib/KernelRunner
WillFroom Jan 31, 2025
6d79342
Integrate LLVM at llvm/llvm-project@4573c857da88
gribozavr Jan 31, 2025
de4d7d3
[xla:gpu] simplify sparsity patches
vwbaker Jan 31, 2025
9e05aac
[XLA:CPU] Save compiled symbols into the cpu executable proto
tensorflower-gardener Jan 31, 2025
8feab2a
[XLA:GPU] Extract more fragments into the smaller functions.
loislo Jan 31, 2025
ca7ff3c
Mark uses of `-std=c++17` in BUILD files as being a workaround for ol…
tensorflower-gardener Jan 31, 2025
7c3720d
[XLA:CPU][roll forward] Underlying ObjectLoader dylibs are using Defi…
tensorflower-gardener Jan 31, 2025
cd8872e
[XLA:GPU] extract 3 different cases from the AddDimToTensorParams
loislo Jan 31, 2025
f45c97c
PR #22102: [ROCm] Add support for the collective memory allocator type
alekstheod Jan 31, 2025
b77fc61
Build IFRT shardings with both addressable and non-addressable device…
emilyfertig Jan 31, 2025
c8bcda1
[xla:cpu] Move ObjectPool and ParallelLoopRunner to top level runtime…
ezhulenev Jan 31, 2025
8b49e39
Merge pull request #63959 from tensorflow:LakshmiKalaKadali-patch-3
tensorflower-gardener Jan 31, 2025
ee7985f
[xla:cpu:xnn] Add microbenchmarks for XNN fusions.
penpornk Jan 31, 2025
7d375b2
[xla:hlo] Always print backend_config when deduplicating computations.
danielsuo Jan 31, 2025
f8049c4
[XLA:CPU] Add debug flag to enable loop unrolling in LLVM IR optimiza…
WillFroom Jan 31, 2025
2eb69d1
PR #21825: Exclude the usage of CPU memory from the GPU memory scheduler
zhenying-liu Jan 31, 2025
5e3baa7
Clean-up LiteRT Compiler error message.
vamsimanchala Jan 31, 2025
330912d
PR #21764: [ROCm] Make hipfft bazel rule publicly visible
hsharsha Jan 31, 2025
bbd9127
[ODML] Pass expand-tuple : Migrate from MHLO to StableHLO
abhigunj Jan 31, 2025
64227f1
Set absolute compiler path for hermetic CUDA repository.
tensorflower-gardener Jan 31, 2025
581e187
Add a possibility to override HostStream by using custom HostStream f…
tensorflower-gardener Jan 31, 2025
190d35a
Display the flax_2b E2E benchmark results to show TTFT and E2E latency
juliagmt-google Jan 31, 2025
44624ab
No public description
tensorflower-gardener Jan 31, 2025
436988e
[xla:cpu] Do not put moved-from XnnRuntime back into the pool
ezhulenev Jan 31, 2025
5bdfe9f
Replace LITERT_ASSERT_STATUS_OK with LITERT_ASSERT_OK.
qukhan Jan 31, 2025
814b7dc
Reverts b77fc611f78c5511c3881060b5c695df9c72d3a8
emilyfertig Jan 31, 2025
350965b
[pjrt] Use the `PjRtMemorySpace` version of `BufferFromHostLiteral`
superbobry Jan 31, 2025
d693e45
PR #22107: Change cuda_configure.bzl to handle blackwell ptx variants…
dimvar Jan 31, 2025
29287e4
Remove redundant TODO, the proposed change has been implemented.
rhundt Jan 31, 2025
7b452d5
Integrate LLVM at llvm/llvm-project@956c0707d909
slackito Jan 31, 2025
93ef858
Update users of TSL headers and targets to new location in XLA
ddunl Jan 31, 2025
9d66452
Reverts 46f98fe577a1591a800d17cbb0c0b4000ff13f74
dimitar-asenov Jan 31, 2025
6f97223
[HLO] Use llvm::StringRef when building MHLO string attributes instea…
GleasonK Jan 31, 2025
135a67d
Reverts 0e749daf4bcf5423c48534611927dd28e985dccc
emilyfertig Jan 31, 2025
2a268f3
Enforece sequential order in send/recv pipeline parallelism test
frgossen Jan 31, 2025
bb8288b
Update users of TSL headers and targets to new location in XLA
ddunl Jan 31, 2025
0e564c9
PR #22073: Add a README file in the gpu_specs directory.
dimvar Jan 31, 2025
73fe797
Remove helper methods from tests to make them easier to read.
toli-y Jan 31, 2025
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
The table of contents is too big for display.
Diff view
Diff view
  •  
  •  
  •  
The diff you're trying to view is too large. We only load the first 3000 changed files.
160 changes: 73 additions & 87 deletions .bazelrc

Large diffs are not rendered by default.

16 changes: 16 additions & 0 deletions .github/ISSUE_TEMPLATE/config.yml
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================

blank_issues_enabled: false
6 changes: 3 additions & 3 deletions CONTRIBUTING.md
Original file line number Diff line number Diff line change
Expand Up @@ -250,7 +250,7 @@ There are two ways to run TensorFlow unit tests.
bazel by doing as follows:

```bash
export flags="--config=opt -k"
export flags="--config=linux -k"
```

If the tests are to be run on the GPU:
Expand All @@ -259,15 +259,15 @@ There are two ways to run TensorFlow unit tests.
flag.

```bash
export flags="--config=opt --config=cuda -k"
export flags="--config=linux --config=cuda -k"
```

* For TensorFlow versions prior v.2.18.0: Add CUDA paths to
LD_LIBRARY_PATH and add the `cuda` option flag.

```bash
export LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64:$LD_LIBRARY_PATH"
export flags="--config=opt --config=cuda -k"
export flags="--config=linux --config=cuda -k"
```

For example, to run all tests under tensorflow/python, do:
Expand Down
2 changes: 1 addition & 1 deletion RELEASE.md
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@
while preserving the implementation flexibility to change the values of
these constants in the future.)
* Interpreter:
* `tf.lite.Interpreter` gives warning of future deletion and a redirection notice to its new location at `ai_edge_litert.interpreter`. See the [migration guide](https://ai.google.dev/edge/litert/migration) for details.
* `tf.lite.Interpreter` gives deprecation warning redirecting to its new location at `ai_edge_litert.interpreter`, as the API `tf.lite.Interpreter` will be deleted in TF 2.20. See the [migration guide](https://ai.google.dev/edge/litert/migration) for details.

### Known Caveats

Expand Down
19 changes: 19 additions & 0 deletions ci/devinfra/docker/windows/Dockerfile
Original file line number Diff line number Diff line change
@@ -1,3 +1,7 @@
# NOTE: This Dockerfile is no longer in use.
# It is kept just in case, but it's recommended to use the 2022 version,
# and that is what internal CI uses as well.

# This Dockerfile creates an image that has:
# - the correct MTU setting for networking from inside the container to work.
# - Visual Studio 2022 Build Tools
Expand Down Expand Up @@ -42,6 +46,7 @@ RUN C:\TEMP\vs_community.exe \
--add Microsoft.VisualStudio.Workload.NativeDesktop \
--add Microsoft.VisualStudio.Component.VC.14.39.17.9.x86.64 \
--add Microsoft.VisualStudio.Component.Windows11SDK.22621 \
--add Microsoft.VisualStudio.Component.VC.ATL \
|| IF "%ERRORLEVEL%"=="3010" EXIT 0

SHELL ["powershell.exe", "-ExecutionPolicy", "Bypass", "-Command", \
Expand Down Expand Up @@ -152,4 +157,18 @@ RUN (New-Object Net.WebClient).DownloadFile( \
$env:PATH = [Environment]::GetEnvironmentVariable('PATH', 'Machine') + ';C:\tools\bazel'; \
[Environment]::SetEnvironmentVariable('PATH', $env:PATH, 'Machine');

ENV CLOUDSDK_CORE_DISABLE_PROMPTS 1
RUN (New-Object Net.WebClient).DownloadFile('https://dl.google.com/dl/cloudsdk/channels/rapid/google-cloud-sdk.zip', 'C:\Temp\google-cloud-sdk.zip'); \
Expand-Archive -Path 'C:\Temp\google-cloud-sdk.zip' -DestinationPath $env:ProgramFiles -Verbose:$false
RUN & \"$env:ProgramFiles\\google-cloud-sdk\\install.bat\" --path-update false
RUN $env:Path += \";$env:ProgramFiles\\google-cloud-sdk\\bin\"; \
[Environment]::SetEnvironmentVariable('Path', $env:Path, [EnvironmentVariableTarget]::Machine);
# Re-enable prompts for interactive use.
ENV CLOUDSDK_CORE_DISABLE_PROMPTS=""

# MSYS attempts to use non-cmd versions, which aren't meant for Windows
RUN Add-Content -Path C:\tools\msys64\.bashrc -Value 'alias gcloud=gcloud.cmd'
RUN Add-Content -Path C:\tools\msys64\.bashrc -Value 'alias gsutil=gsutil.cmd'
RUN Add-Content -Path C:\tools\msys64\.bashrc -Value 'alias bq=bq.cmd'

SHELL ["cmd.exe", "/s", "/c"]
207 changes: 207 additions & 0 deletions ci/devinfra/docker/windows2022/Dockerfile
Original file line number Diff line number Diff line change
@@ -0,0 +1,207 @@
# This Dockerfile creates an image that has:
# - the correct MTU setting for networking from inside the container to work.
# - Visual Studio 2022 Build Tools
# - MSVC 14.39
# - LLVM/Clang 18.1.4
# - MSYS2 + curl, git, patch, vim, unzip, zip
# - Python 3.9 - 3.13
# - Bazelisk 1.22.1
# - JDK 21 (Azul Zulu)

FROM mcr.microsoft.com/windows/servercore:ltsc2022

SHELL ["powershell.exe", "-ExecutionPolicy", "Bypass", "-Command", \
"$ErrorActionPreference='Stop'; $ProgressPreference='SilentlyContinue';$VerbosePreference = 'Continue';"]

# Enable long paths
RUN New-ItemProperty -Path "HKLM:\SYSTEM\CurrentControlSet\Control\FileSystem" -Name "LongPathsEnabled" -Value 1 -PropertyType DWORD -Force

RUN md C:\TEMP
RUN md C:\TMP
ENV TMP "C:/TMP"
ENV TEMP "C:/TEMP"

# Install 7-Zip.
RUN (New-Object Net.WebClient).DownloadFile('https://www.7-zip.org/a/7z2201-x64.msi', '7z.msi'); \
Start-Process msiexec.exe -ArgumentList \"/i 7z.msi /qn /norestart /log C:\\TEMP\\7z_install_log.txt\" -wait; \
Remove-Item .\7z.msi;

# Download the Visual Studio 2022 Installer.
RUN (New-Object Net.WebClient).DownloadFile('https://aka.ms/vs/17/release/vs_community.exe', 'C:\TEMP\vs_community.exe');
# Install Visual Studio 2022 Build Tools + Compiler
SHELL ["cmd", "/S", "/C"]
# Packages, and component versions, can be found here:
# https://learn.microsoft.com/en-us/visualstudio/install/workload-component-id-vs-build-tools
RUN C:\TEMP\vs_community.exe \
--quiet --wait --norestart --nocache \
--add Microsoft.VisualStudio.Component.VC.Tools.x86.x64 \
--add Microsoft.VisualStudio.Workload.NativeDesktop \
--add Microsoft.VisualStudio.Component.VC.14.39.17.9.x86.64 \
--add Microsoft.VisualStudio.Component.Windows11SDK.22621 \
|| IF "%ERRORLEVEL%"=="3010" EXIT 0

SHELL ["powershell.exe", "-ExecutionPolicy", "Bypass", "-Command", \
"$ErrorActionPreference='Stop'; $ProgressPreference='SilentlyContinue'; $VerbosePreference = 'Continue';"]

# Install Clang.
RUN (New-Object Net.WebClient).DownloadFile( \
'https://github.com/llvm/llvm-project/releases/download/llvmorg-18.1.4/LLVM-18.1.4-win64.exe', \
'LLVM.exe'); \
Start-Process -FilePath \"C:\Program Files\7-Zip\7z.exe\" -ArgumentList 'x LLVM.exe -oC:\tools\LLVM' -Wait; \
$env:PATH = [Environment]::GetEnvironmentVariable('PATH', 'Machine') + ';C:\tools\LLVM\bin'; \
[Environment]::SetEnvironmentVariable('PATH', $env:PATH, 'Machine');

# Install MSYS2.
RUN (New-Object Net.WebClient).DownloadFile( \
'https://repo.msys2.org/distrib/x86_64/msys2-base-x86_64-20240727.tar.xz', \
'msys2.tar.xz'); \
Start-Process -FilePath \"C:\Program Files\7-Zip\7z.exe\" -ArgumentList 'x msys2.tar.xz -oC:\TEMP\msys2.tar' -Wait; \
Start-Process -FilePath \"C:\Program Files\7-Zip\7z.exe\" -ArgumentList 'x C:\TEMP\msys2.tar -oC:\tools' -Wait; \
$env:PATH = [Environment]::GetEnvironmentVariable('PATH', 'Machine') + ';C:\tools\msys64;C:\tools\msys64\usr\bin\'; \
[Environment]::SetEnvironmentVariable('PATH', $env:PATH, 'Machine');

# Disable signature checking on pacman because we cannot initialize the keyring.
RUN Add-Content -Path C:\tools\msys64\etc\pacman.d\mirrorlist.mingw32 -Value 'SigLevel = Never'
RUN Add-Content -Path C:\tools\msys64\etc\pacman.d\mirrorlist.mingw64 -Value 'SigLevel = Never'
RUN Add-Content -Path C:\tools\msys64\etc\pacman.d\mirrorlist.msys -Value 'SigLevel = Never'

# Install pacman packages.
RUN C:\tools\msys64\usr\bin\bash.exe -lc \
'pacman --noconfirm -Syy curl git patch vim unzip zip'

# Install multiple Pythons, but only add one of them to PATH.
RUN function Install-Python { \
param( \
[string]$version, \
[int]$prependPath \
) \
$url = ('https://www.python.org/ftp/python/{0}/python-{0}-amd64.exe' -f $version); \
Write-Host ('Downloading {0}...' -f $url); \
[Net.ServicePointManager]::SecurityProtocol = [Net.SecurityProtocolType]::Tls12; \
(New-Object Net.WebClient).DownloadFile($url, 'C:\tmp\pyinstall.exe'); \
\
# Without the patch version \
$truncatedVersion = $($version -replace '\.\d+$', ''); \
$installDir = 'C:\Python' + $truncatedVersion; \
Write-Host ('Installing into {0} (PrependPath: {1})...' -f $installDir, $($prependPath -eq 1)); \
$argumentList = ('/quiet InstallAllUsers=1 PrependPath={0} TargetDir={1}' -f $prependPath, $installDir); \
Start-Process -FilePath 'C:\tmp\pyinstall.exe' -ArgumentList $argumentList -Wait; \
\
Write-Host 'Verifying install...'; \
Write-Host \" python --version $version\"; & $installDir\python.exe --version; \
\
Write-Host 'Verifying pip install...'; \
& $installDir\python.exe -m pip --version; \
\
Write-Host 'Updating pip...'; \
& $installDir\python.exe -m pip install --upgrade pip; \
\
Write-Host 'Installing/updating packages...'; \
& $installDir\python.exe -m pip install --upgrade setuptools packaging; \
\
Write-Host 'Removing installation binary...'; \
Remove-Item C:\tmp\pyinstall.exe -Force; \
}; \
Write-Host 'Installing multiple Python versions...'; \
$versions = @( \
@{ version = '3.9.13'; prependPath = 0 }, \
@{ version = '3.10.11'; prependPath = 0 }, \
@{ version = '3.11.9'; prependPath = 0 }, \
@{ version = '3.12.8'; prependPath = 0 }, \
@{ version = '3.13.1'; prependPath = 1 } \
); \
foreach ($v in $versions) { \
Install-Python -version $v.version -prependPath $v.prependPath; \
}; \
Write-Host 'Python installations complete.';

# Add a python3 symlink for the Python in PATH.
# It's not feasible to add one for each version, as on Windows,
# Python uses PATH and the binary's (symlink's) location, to launch itself.
RUN C:\tools\msys64\usr\bin\bash.exe -lc 'ln -s /c/Python3.13/python.exe /usr/bin/python3';

# Install JDK 21.
RUN \
Add-Type -AssemblyName \"System.IO.Compression.FileSystem\"; \
$zulu_pkg = \"zulu21.34.19-ca-jdk21.0.3-win_x64.zip\"; \
$zulu_url = \"https://cdn.azul.com/zulu/bin/${zulu_pkg}\"; \
$zulu_zip = \"c:\\temp\\${zulu_pkg}\"; \
$zulu_extracted_path = \"c:\\temp\\\" + [IO.Path]::GetFileNameWithoutExtension($zulu_zip); \
$zulu_root = \"c:\\openjdk\"; \
(New-Object Net.WebClient).DownloadFile($zulu_url, $zulu_zip); \
[System.IO.Compression.ZipFile]::ExtractToDirectory($zulu_zip, \"c:\\temp\"); \
Move-Item $zulu_extracted_path -Destination $zulu_root; \
Remove-Item $zulu_zip; \
$env:PATH = [Environment]::GetEnvironmentVariable(\"PATH\", \"Machine\") + \";${zulu_root}\\bin\"; \
[Environment]::SetEnvironmentVariable(\"PATH\", $env:PATH, \"Machine\"); \
$env:JAVA_HOME = $zulu_root; \
[Environment]::SetEnvironmentVariable(\"JAVA_HOME\", $env:JAVA_HOME, \"Machine\")

# Point to the LLVM installation.
# The Bazel Windows guide claims it can find LLVM automatically,
# but it likely only works if it's installed somewhere inside C:\Program Files.
ENV BAZEL_LLVM "C:\tools\LLVM"

# These variables may be useful, but so far haven't been. Keeping for posterity.
# ENV CLANG_COMPILER_PATH "C:\tools\llvm\bin\clang.exe"
# ENV CC "C:\tools\llvm\bin\clang.exe"
# ENV BAZEL_COMPILER "C:\tools\llvm\bin\clang.exe"

ENV BAZEL_SH "C:\tools\msys64\usr\bin\bash.exe"
ENV BAZEL_VS "C:\Program Files\Microsoft Visual Studio\2022\BuildTools"
ENV BAZEL_VC "C:\Program Files\Microsoft Visual Studio\2022\Community\VC"

# Environment variables to prevent auto-conversion of Linux-like paths to Windows paths.
# This is necessary as some paths end up invalid/mangled.
ENV MSYS_NO_PATHCONV 1
ENV MSYS2_ARG_CONV_EXCL *

# This should only be necessary if there are multiple, differently-versioned
# MSVC compilers installed, and a particular one should be used.
# To find exact versions available:
# - Navigate to the relevant folder, e.g.
# C:\Program Files\Microsoft Visual Studio\2022
# - Search for the `cl.exe` file: `gci -r -fi cl.exe`
# - The version will be part of the found path, e.g.
# 2022\Community\VC\Tools\MSVC\14.39.33519\bin\Hostx64\x64
# ENV BAZEL_VC_FULL_VERSION 14.39.33519

# Install Bazelisk.
RUN md C:\tools\bazel
RUN (New-Object Net.WebClient).DownloadFile( \
'https://github.com/bazelbuild/bazelisk/releases/download/v1.22.1/bazelisk-windows-amd64.exe', \
'C:\tools\bazel\bazel.exe'); \
$env:PATH = [Environment]::GetEnvironmentVariable('PATH', 'Machine') + ';C:\tools\bazel'; \
[Environment]::SetEnvironmentVariable('PATH', $env:PATH, 'Machine');

# Install gcloud, and add it to PATH
ENV CLOUDSDK_CORE_DISABLE_PROMPTS 1
RUN (New-Object Net.WebClient).DownloadFile('https://dl.google.com/dl/cloudsdk/channels/rapid/google-cloud-sdk.zip', 'C:\Temp\google-cloud-sdk.zip'); \
Expand-Archive -Path 'C:\Temp\google-cloud-sdk.zip' -DestinationPath $env:ProgramFiles -Verbose:$false
RUN & \"$env:ProgramFiles\\google-cloud-sdk\\install.bat\" --path-update false
RUN $env:Path += \";$env:ProgramFiles\\google-cloud-sdk\\bin\"; \
[Environment]::SetEnvironmentVariable('Path', $env:Path, [EnvironmentVariableTarget]::Machine);
# Re-enable prompts for interactive use.
ENV CLOUDSDK_CORE_DISABLE_PROMPTS=""

# MSYS attempts to use non-cmd versions, which aren't meant for Windows
RUN Add-Content -Path C:\tools\msys64\.bashrc -Value 'alias gcloud=gcloud.cmd'
RUN Add-Content -Path C:\tools\msys64\.bashrc -Value 'alias gsutil=gsutil.cmd'
RUN Add-Content -Path C:\tools\msys64\.bashrc -Value 'alias bq=bq.cmd'

# Symlink a directory, to have it pretend be the T:\ drive.
# This drive letter is used by internal CI,
# and part of it is mounted to the container during the container's creation.
#
# While the mount argument (`-v host_path:container_path`) still requires
# `container_path` to be a legitimate C:\ path, in this case, 'C:\drive_t',
# this symlink does allow for the convenience of passing unedited paths
# to `docker exec` commands, e.g., 'T:\path', instead of 'C:\path',
# without having to replace the drive letter with C:\ every time.
# Such a workaround is not required on Linux, since it
# can create arbitrary paths within the container, e.g., '/t'.
# Note: This does not affect/work for `docker cp` commands.
RUN New-Item -ItemType directory -Path C:\drive_t; \
New-ItemProperty -Path 'HKLM:\SYSTEM\CurrentControlSet\Control\Session Manager\DOS Devices' -Name 'T:' -Value '\??\C:\drive_t' -PropertyType String;

CMD ["bash", "-i"]
2 changes: 1 addition & 1 deletion ci/official/any.sh
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@
# export TF_ANY_EXTRA_ENV=ci/official/envs/local_rbe
# ./any.sh
# ...
set -euxo pipefail
set -exo pipefail
cd "$(dirname "$0")/../../" # tensorflow/
# Any request that includes "nightly_upload" should just use the
# local multi-cache (public read-only cache + disk cache) instead.
Expand Down
2 changes: 1 addition & 1 deletion ci/official/bisect.sh
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@
# export TF_BISECT_BAD=a_failing_commit_sha
# export TF_ANY_TARGETS="quoted list of targets, like on the command line"
# export TF_ANY_MODE=test
set -euxo pipefail
set -exo pipefail
cd "$(dirname "$0")/../../" # tensorflow/
export TFCI="$(echo $TFCI | sed 's/,nightly_upload/,public_cache,disk_cache/')"
git bisect start "$TF_BISECT_BAD" "$TF_BISECT_GOOD"
Expand Down
2 changes: 2 additions & 0 deletions ci/official/containers/linux_arm64/Dockerfile
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,8 @@ RUN /setup.python.sh python3.10 devel.requirements.txt
RUN /setup.python.sh python3.11 devel.requirements.txt
RUN /setup.python.sh python3.12 devel.requirements.txt
RUN /setup.python.sh python3.13 devel.requirements.txt
# python3.13-nogil is a free-threaded build of python3.13
RUN /setup.python.sh python3.13-nogil devel.requirements.txt

FROM devel as tf
# Setup TF Python environment.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@

# Check and rename wheels with auditwheel. Inserts the platform tags like
# "manylinux_xyz" into the wheel filename.
set -euxo pipefail
set -exo pipefail

for wheel in /tf/pkg/*.whl; do
echo "Checking and renaming $wheel..."
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
set -euxo pipefail
set -exo pipefail

# Run this from inside the tensorflow github directory.
# Usage: setup_venv_test.sh venv_and_symlink_name "glob pattern for one wheel file"
Expand Down
4 changes: 1 addition & 3 deletions ci/official/containers/linux_arm64/setup.python.sh
Original file line number Diff line number Diff line change
Expand Up @@ -25,11 +25,9 @@ REQUIREMENTS=$2

add-apt-repository ppa:deadsnakes/ppa
# Install Python packages for this container's version
if [[ ${VERSION} == "python3.13" ]]; then
if [[ ${VERSION} == "python3.13-nogil" ]]; then
cat >pythons.txt <<EOF
$VERSION
$VERSION-dev
$VERSION-venv
EOF
else
cat >pythons.txt <<EOF
Expand Down
18 changes: 12 additions & 6 deletions ci/official/containers/ml_build/Dockerfile
Original file line number Diff line number Diff line change
@@ -1,5 +1,8 @@
################################################################################
FROM ubuntu:22.04@sha256:58b87898e82351c6cf9cf5b9f3c20257bb9e2dcf33af051e12ce532d7f94e3fe AS devel
ARG BASE_IMAGE=ubuntu:22.04@sha256:58b87898e82351c6cf9cf5b9f3c20257bb9e2dcf33af051e12ce532d7f94e3fe
FROM $BASE_IMAGE AS devel
# See https://docs.docker.com/reference/dockerfile/#understand-how-arg-and-from-interact
# on why we cannot reference BASE_IMAGE again unless we declare it again.
################################################################################

# Install devtoolset build dependencies
Expand All @@ -20,15 +23,15 @@ RUN /build_devtoolset.sh devtoolset-9 /dt9
# Setup Python
COPY setup.python.sh /setup.python.sh
COPY builder.requirements.txt /builder.requirements.txt
RUN /setup.python.sh python3.9 builder.requirements.txt
RUN /setup.python.sh python3.10 builder.requirements.txt
RUN /setup.python.sh python3.11 builder.requirements.txt
RUN /setup.python.sh python3.13 builder.requirements.txt
RUN /setup.python.sh python3.9 /builder.requirements.txt
RUN /setup.python.sh python3.10 /builder.requirements.txt
RUN /setup.python.sh python3.11 /builder.requirements.txt
RUN /setup.python.sh python3.13 /builder.requirements.txt

# Since we are using python3.12 as the default python version, we need to
# install python3.12 last for now.
# TODO(b/376338367): switch to pyenv.
RUN /setup.python.sh python3.12 builder.requirements.txt
RUN /setup.python.sh python3.12 /builder.requirements.txt

# Setup links for TensorFlow to compile.
# Referenced in devel.usertools/*.bazelrc.
Expand All @@ -41,6 +44,9 @@ RUN ln -sf /usr/lib/python3.12 /usr/lib/tf_python
# Make sure clang is on the path
RUN ln -s /usr/lib/llvm-18/bin/clang /usr/bin/clang

# Link the compat driver to the location if available.
RUN if [ -e "/usr/local/cuda/compat/libcuda.so.1" ]; then ln -s /usr/local/cuda/compat/libcuda.so.1 /usr/lib/x86_64-linux-gnu/libcuda.so.1; fi

# Install various tools.
# - bats: bash unit testing framework
# - bazelisk: always use the correct bazel version
Expand Down
Loading