Skip to content

Commit 073bc92

Browse files
github-actions[bot]github-actions[bot]
authored
Merge pull request #227 from ROCm/ci-upstream-sync-110_1
CI: 02/10/25 upstream sync
2 parents c501484 + 9133253 commit 073bc92

File tree

181 files changed

+7971
-5204
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

181 files changed

+7971
-5204
lines changed

‎.github/workflows/bazel_cpu_rbe.yml‎

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,10 +14,15 @@ on:
1414
pull_request:
1515
branches:
1616
- main
17+
push:
18+
branches:
19+
- main
20+
- 'release/**'
1721

1822
concurrency:
1923
group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
20-
cancel-in-progress: true
24+
# Don't cancel in-progress jobs for main/release branches.
25+
cancel-in-progress: ${{ !contains(github.ref, 'release/') && github.ref != 'main' }}
2126

2227
jobs:
2328
run_tests:
@@ -26,14 +31,22 @@ jobs:
2631
container: ${{ (contains(matrix.runner, 'linux-x86') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest') ||
2732
(contains(matrix.runner, 'linux-arm64') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build-arm64:latest') }}
2833
env:
29-
JAXCI_HERMETIC_PYTHON_VERSION: "3.12"
34+
JAXCI_HERMETIC_PYTHON_VERSION: ${{ matrix.python }}
3035
JAXCI_ENABLE_X64: ${{ matrix.enable-x_64 }}
3136
# Begin Presubmit Naming Check - name modification requires internal check to be updated
3237
strategy:
3338
matrix:
39+
python: ["3.10", "3.13"]
3440
runner: ["linux-x86-n2-16", "linux-arm64-c4a-16"]
3541
enable-x_64: [1, 0]
36-
name: "Bazel CPU tests (${{ matrix.runner }}, Python 3.12, x64=${{ matrix.enable-x_64 }})"
42+
exclude:
43+
# Exclude x64=1 on the oldest Python and x64=0 on the newest Python. As long as we have
44+
# coverage for one of each, we don't need to run both.
45+
- python: "3.10"
46+
enable-x_64: 1
47+
- python: "3.13"
48+
enable-x_64: 0
49+
name: "Bazel CPU tests (${{ matrix.runner }}, Python ${{ matrix.python }}, x64=${{ matrix.enable-x_64 }})"
3750
# End Presubmit Naming Check github-cpu-presubmits
3851
steps:
3952
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2

‎.github/workflows/bazel_cuda_rbe.yml‎

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,25 +14,38 @@ on:
1414
pull_request:
1515
branches:
1616
- main
17+
push:
18+
branches:
19+
- main
20+
- 'release/**'
1721

1822
concurrency:
1923
group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
20-
cancel-in-progress: true
24+
# Don't cancel in-progress jobs for main/release branches.
25+
cancel-in-progress: ${{ !contains(github.ref, 'release/') && github.ref != 'main' }}
2126

2227
jobs:
2328
run_tests:
2429
if: github.event.repository.fork == false
2530
runs-on: ${{ matrix.runner }}
2631
container: 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest'
2732
env:
28-
JAXCI_HERMETIC_PYTHON_VERSION: "3.12"
33+
JAXCI_HERMETIC_PYTHON_VERSION: ${{ matrix.python }}
2934
JAXCI_ENABLE_X64: ${{ matrix.enable-x_64 }}
3035
# Begin Presubmit Naming Check - name modification requires internal check to be updated
3136
strategy:
3237
matrix:
38+
python: ["3.10", "3.13"]
3339
runner: ["linux-x86-n2-16"]
3440
enable-x_64: [1, 0]
35-
name: "Bazel single accelerator CUDA tests (${{ matrix.runner }}, Python 3.12, x64=${{ matrix.enable-x_64 }})"
41+
exclude:
42+
# Exclude x64=1 on the oldest Python and x64=0 on the newest Python. As long as we have
43+
# coverage for one of each, we don't need to run both.
44+
- python: "3.10"
45+
enable-x_64: 1
46+
- python: "3.13"
47+
enable-x_64: 0
48+
name: "Bazel single accelerator CUDA tests (${{ matrix.runner }}, Python ${{ matrix.python }}, x64=${{ matrix.enable-x_64 }})"
3649
# End Presubmit Naming Check github-cuda-presubmits
3750
steps:
3851
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2

‎.github/workflows/ci-build.yaml‎

Lines changed: 19 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ jobs:
3131
steps:
3232
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
3333
- name: Set up Python 3.11
34-
uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0
34+
uses: actions/setup-python@42375524e23c412d93fb67b49958b491fce71c38 # v5.4.0
3535
with:
3636
python-version: 3.11
3737
- run: python -m pip install pre-commit
@@ -62,22 +62,13 @@ jobs:
6262
steps:
6363
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
6464
- name: Set up Python ${{ matrix.python-version }}
65-
uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0
65+
uses: actions/setup-python@42375524e23c412d93fb67b49958b491fce71c38 # v5.4.0
6666
with:
6767
python-version: ${{ matrix.python-version }}
68-
- name: Get pip cache dir
69-
id: pip-cache
70-
run: |
71-
python -m pip install --upgrade pip wheel
72-
echo "dir=$(pip cache dir)" >> $GITHUB_OUTPUT
73-
- name: pip cache
74-
uses: actions/cache@1bd1e32a3bdc45362d1e726936510720a7c30a57 # v4.2.0
75-
with:
76-
path: ${{ steps.pip-cache.outputs.dir }}
77-
key: ${{ runner.os }}-py${{ matrix.python-version }}-pip-${{ hashFiles('**/setup.py', '**/requirements.txt', '**/test-requirements.txt') }}
7868
- name: Install dependencies
7969
run: |
80-
pip install .[minimum-jaxlib] -r build/test-requirements.txt
70+
pip install uv
71+
uv pip install --system .[minimum-jaxlib] -r build/test-requirements.txt
8172
8273
- name: Run tests
8374
env:
@@ -109,30 +100,21 @@ jobs:
109100
steps:
110101
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
111102
- name: Set up Python ${{ matrix.python-version }}
112-
uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0
103+
uses: actions/setup-python@42375524e23c412d93fb67b49958b491fce71c38 # v5.4.0
113104
with:
114105
python-version: ${{ matrix.python-version }}
115-
- name: Get pip cache dir
116-
id: pip-cache
117-
run: |
118-
python -m pip install --upgrade pip wheel
119-
echo "dir=$(pip cache dir)" >> $GITHUB_OUTPUT
120-
- name: pip cache
121-
uses: actions/cache@1bd1e32a3bdc45362d1e726936510720a7c30a57 # v4.2.0
122-
with:
123-
path: ${{ steps.pip-cache.outputs.dir }}
124-
key: ${{ runner.os }}-pip-docs-${{ hashFiles('**/setup.py', '**/requirements.txt', '**/test-requirements.txt') }}
125106
- name: Install dependencies
126107
run: |
127-
pip install -r docs/requirements.txt
108+
pip install uv
109+
uv pip install --system -r docs/requirements.txt
128110
- name: Test documentation
129111
env:
130112
XLA_FLAGS: "--xla_force_host_platform_device_count=8"
131113
JAX_TRACEBACK_FILTERING: "off"
132114
JAX_ARRAY: 1
133115
PY_COLORS: 1
134116
run: |
135-
pytest -n auto --tb=short --doctest-glob='*.md' --doctest-glob='*.rst' docs --doctest-continue-on-failure --ignore=docs/multi_process.md
117+
pytest -n auto --tb=short --doctest-glob='*.md' --doctest-glob='*.rst' docs --doctest-continue-on-failure --ignore=docs/multi_process.md
136118
pytest -n auto --tb=short --doctest-modules jax --ignore=jax/config.py --ignore=jax/experimental/jax2tf --ignore=jax/_src/lib/mlir --ignore=jax/_src/lib/triton.py --ignore=jax/_src/lib/mosaic_gpu.py --ignore=jax/interpreters/mlir.py --ignore=jax/experimental/array_serialization --ignore=jax/collect_profile.py --ignore=jax/_src/tpu_custom_call.py --ignore=jax/experimental/mosaic --ignore=jax/experimental/pallas --ignore=jax/_src/pallas --ignore=jax/lib/xla_extension.py
137119
138120
@@ -146,22 +128,13 @@ jobs:
146128
steps:
147129
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
148130
- name: Set up Python ${{ matrix.python-version }}
149-
uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0
131+
uses: actions/setup-python@42375524e23c412d93fb67b49958b491fce71c38 # v5.4.0
150132
with:
151133
python-version: ${{ matrix.python-version }}
152-
- name: Get pip cache dir
153-
id: pip-cache
154-
run: |
155-
python -m pip install --upgrade pip wheel
156-
echo "dir=$(pip cache dir)" >> $GITHUB_OUTPUT
157-
- name: pip cache
158-
uses: actions/cache@1bd1e32a3bdc45362d1e726936510720a7c30a57 # v4.2.0
159-
with:
160-
path: ${{ steps.pip-cache.outputs.dir }}
161-
key: ${{ runner.os }}-pip-docs-${{ hashFiles('**/setup.py', '**/requirements.txt', '**/test-requirements.txt') }}
162134
- name: Install dependencies
163135
run: |
164-
pip install -r docs/requirements.txt
136+
pip install uv
137+
uv pip install --system -r docs/requirements.txt
165138
- name: Render documentation
166139
run: |
167140
sphinx-build -j auto --color -W --keep-going -b html -D nb_execution_mode=off docs docs/build/html
@@ -181,22 +154,13 @@ jobs:
181154
steps:
182155
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
183156
- name: Set up Python ${{ matrix.python-version }}
184-
uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0
157+
uses: actions/setup-python@42375524e23c412d93fb67b49958b491fce71c38 # v5.4.0
185158
with:
186159
python-version: ${{ matrix.python-version }}
187-
- name: Get pip cache dir
188-
id: pip-cache
189-
run: |
190-
python -m pip install --upgrade pip wheel
191-
echo "dir=$(pip cache dir)" >> $GITHUB_OUTPUT
192-
- name: pip cache
193-
uses: actions/cache@1bd1e32a3bdc45362d1e726936510720a7c30a57 # v4.2.0
194-
with:
195-
path: ${{ steps.pip-cache.outputs.dir }}
196-
key: ${{ runner.os }}-py${{ matrix.python-version }}-pip-${{ hashFiles('**/setup.py', '**/requirements.txt', '**/test-requirements.txt') }}
197160
- name: Install dependencies
198161
run: |
199-
pip install .[minimum-jaxlib] tensorflow -r build/test-requirements.txt
162+
pip install uv
163+
uv pip install --system .[minimum-jaxlib] tensorflow -r build/test-requirements.txt
200164
201165
- name: Run tests
202166
env:
@@ -220,23 +184,15 @@ jobs:
220184
steps:
221185
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
222186
- name: Set up Python
223-
uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0
187+
uses: actions/setup-python@42375524e23c412d93fb67b49958b491fce71c38 # v5.4.0
224188
with:
225189
python-version: 3.12
226-
- name: Get pip cache dir
227-
id: pip-cache
228-
run: |
229-
python -m pip install --upgrade pip wheel
230-
echo "dir=$(pip cache dir)" >> $GITHUB_OUTPUT
231-
- name: pip cache
232-
uses: actions/cache@1bd1e32a3bdc45362d1e726936510720a7c30a57 # v4.2.0
233-
with:
234-
path: ${{ steps.pip-cache.outputs.dir }}
235-
key: ${{ runner.os }}-pip-ffi-examples-${{ hashFiles('**/setup.py', '**/requirements.txt', '**/test-requirements.txt', 'examples/**/pyproject.toml') }}
236190
- name: Install JAX
237-
run: pip install .
191+
run: |
192+
pip install uv
193+
uv pip install --system .
238194
- name: Build and install example project
239-
run: python -m pip install -v ./examples/ffi[test]
195+
run: uv pip install --system ./examples/ffi[test]
240196
env:
241197
# We test building using GCC instead of clang. All other JAX builds use
242198
# clang, but it is useful to make sure that FFI users can compile using

‎.github/workflows/cloud-tpu-ci-presubmit.yml‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,10 @@ on:
1717
pull_request:
1818
branches:
1919
- main
20+
push:
21+
branches:
22+
- main
23+
- 'release/**'
2024

2125
# This should also be set to read-only in the project settings, but it's nice to
2226
# document and enforce the permissions here.
@@ -25,7 +29,8 @@ permissions:
2529

2630
concurrency:
2731
group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
28-
cancel-in-progress: true
32+
# Don't cancel in-progress jobs for main/release branches.
33+
cancel-in-progress: ${{ !contains(github.ref, 'release/') && github.ref != 'main' }}
2934

3035
jobs:
3136
cloud-tpu-test:

‎.github/workflows/jax-array-api.yml‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -32,13 +32,14 @@ jobs:
3232
submodules: 'true'
3333
path: 'array-api-tests'
3434
- name: Set up Python ${{ matrix.python-version }}
35-
uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0
35+
uses: actions/setup-python@42375524e23c412d93fb67b49958b491fce71c38 # v5.4.0
3636
with:
3737
python-version: ${{ matrix.python-version }}
3838
- name: Install dependencies
3939
run: |
40-
python -m pip install .[ci]
41-
python -m pip install pytest-xdist -r array-api-tests/requirements.txt
40+
pip install uv
41+
uv pip install --system .[ci]
42+
uv pip install --system pytest-xdist -r array-api-tests/requirements.txt
4243
- name: Run the test suite
4344
env:
4445
ARRAY_API_TESTS_MODULE: jax.numpy

‎.github/workflows/tsan.yaml‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ concurrency:
66

77
on:
88
schedule:
9-
- cron: "0 12 * * *" # Daily at 12:00 UTC
9+
- cron: "0 5 * * *" # Daily at 05:00 UTC == 00:00 EST == 21:00 PST
1010
workflow_dispatch: # allows triggering the workflow run manually
1111
pull_request: # Automatically trigger on pull requests affecting this file
1212
branches:
@@ -50,15 +50,15 @@ jobs:
5050
path: numpy
5151
submodules: true
5252

53-
- name: Restore cached CPython with TSAN
53+
- name: Restore cached TSAN CPython
5454
id: cache-cpython-tsan-restore
5555
uses: actions/cache/restore@1bd1e32a3bdc45362d1e726936510720a7c30a57 # v4.2.0
5656
with:
5757
path: |
5858
./python-tsan.tgz
5959
key: ${{ runner.os }}-cpython-tsan-${{ hashFiles('cpython/configure.ac') }}
6060

61-
- name: Build CPython with TSAN enabled
61+
- name: Build CPython with enabled TSAN
6262
if: steps.cache-cpython-tsan-restore.outputs.cache-hit != 'true'
6363
run: |
6464
cd cpython
@@ -173,6 +173,10 @@ jobs:
173173
--bazel_options=--copt=-g \
174174
--clang_path=/usr/bin/clang-18
175175
176+
# Update the patch to use TSAN instrumented numpy
177+
sed -i "s|+--extra-index-url.*|+--extra-index-url file://${GITHUB_WORKSPACE}/wheelhouse/|" .github/workflows/requirements_lock_3_13_ft.patch
178+
cat .github/workflows/requirements_lock_3_13_ft.patch
179+
176180
# Apply a patch to numpy in requirements lock 3.13 ft to use the nightly version
177181
git apply .github/workflows/requirements_lock_3_13_ft.patch
178182
@@ -194,4 +198,8 @@ jobs:
194198
--test_output=errors \
195199
--local_test_jobs=32 \
196200
--test_timeout=600 \
201+
--config=resultstore \
202+
--spawn_strategy=local \
203+
--remote_cache=remotebuildexecution.googleapis.com \
204+
--remote_instance_name=projects/tensorflow-testing/instances/default_instance \
197205
//tests:cpu_tests

‎.github/workflows/upstream-nightly.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ jobs:
3333
steps:
3434
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
3535
- name: Set up Python ${{ matrix.python-version }}
36-
uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0
36+
uses: actions/setup-python@42375524e23c412d93fb67b49958b491fce71c38 # v5.4.0
3737
with:
3838
python-version: ${{ matrix.python-version }}
3939
- name: Install JAX test requirements

‎.github/workflows/wheel_win_x64.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ jobs:
2727

2828
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
2929

30-
- uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0
30+
- uses: actions/setup-python@42375524e23c412d93fb67b49958b491fce71c38 # v5.4.0
3131
with:
3232
python-version: ${{ matrix.pyver }}
3333
cache: 'pip'

‎.github/workflows/windows_ci.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ jobs:
3535
with:
3636
path: jax
3737

38-
- uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0
38+
- uses: actions/setup-python@42375524e23c412d93fb67b49958b491fce71c38 # v5.4.0
3939
with:
4040
python-version: ${{ matrix.pyver }}
4141
cache: 'pip'

‎CHANGELOG.md‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,11 +21,15 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
2121
decorator to support customizing the behavior of opaque functions under
2222
JAX-level dead code elimination (DCE). See {jax-issue}`#25956` for more
2323
details.
24-
* Added {func}`jax.random.multinomial`.
2524

2625
* Changes
2726
* `JAX_CPU_COLLECTIVES_IMPLEMENTATION` and `JAX_NUM_CPU_DEVICES` now work as
2827
env vars. Before they could only be specified via jax.config or flags.
28+
* `JAX_CPU_COLLECTIVES_IMPLEMENTATION` now defaults to `'gloo'`, meaning
29+
multi-process CPU communication works out-of-the-box.
30+
* The `jax[tpu]` TPU extra no longer depends on the `libtpu-nightly` package.
31+
This package may safely be removed if it is present on your machine; JAX now
32+
uses `libtpu` instead.
2933

3034
## jax 0.5.0 (Jan 17, 2025)
3135

0 commit comments

Comments
 (0)