Skip to content

CI: 04/10/25 upstream sync #355

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

Closed
wants to merge 624 commits into from

Conversation

rocm-repo-management-api-2[bot]
Copy link

Daily sync with upstream

mattjj and others added 30 commits April 1, 2025 20:18
PiperOrigin-RevId: 742840628
PiperOrigin-RevId: 742927160
Treat `None` metadata values as a special instruction not to set (or
to unset, if nested) the corresponding entry.

In particular, this makes it possible to unset metadata within the
sub-computations of higher-order operations (e.g. branches in
conditionals, loop bodies, etc.). This can be used, for example, to
annotate a conditional but not all the operations in its
branches. That is, the HLO for the following function `f` on a scalar
float argument:

```
def cos(x):
  with set_xla_metadata(a=None):
    return jnp.cos(x)

@jax.jit
def f(x):
  with set_xla_metadata(a="b"):
    return jax.lax.cond(x < 0., jnp.sin, cos, x)
```

produces an attribute `a` on the conditional and on the sine, but not
on the cosine.
Added tsan ci cpython 3.14 job
kwargs are passed sorted by the actual kwarg keyword. This order
must be accounted for when we construct the `debug_info.arg_names`.

Extended the tests to be more precise about not mixing up kwargs,
e.g., use different shapes and look for the shape in the HLO.
PiperOrigin-RevId: 743074472
…r shardings if they are not empty. The same thing happens in select_p's sharding rule

PiperOrigin-RevId: 743105350
It has been more than 6 months since the release of 0.4.32 which was the first release to stop including `permutation_size` as an attribute when lowering, so it is now safe (via our compatibility policy) to remove this argument.

PiperOrigin-RevId: 743132169
Following the compatibility timeline described here: https://docs.jax.dev/en/latest/export/export.html#ensuring-forward-and-backward-compatibility

On Apr 2, it will have been 6 months since the release of 0.4.34 which is the relevant release for this kernels.

PiperOrigin-RevId: 743142261
Description:
- Disabled default env var LOBPCG_EMIT_DEBUG_PLOTS=1
- When run inside TSAN CI job with 3.14t cpython and under multi-threading
the test code from main leads to `RecursionError: maximum recursion depth exceeded` error:
```
ERROR: testLobpcgMonotonicityF32cluster_k_1__n100 (__main__.F32LobpcgTest)
F32LobpcgTest.testLobpcgMonotonicityF32cluster_k_1__n100
testLobpcgMonotonicityF32cluster_k_1__n100(matrix_name='cluster(k-1)', n=100, k=10, m=20, tol=2e-06)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/root/.cache/bazel/_bazel_root/840503f2165a538d6d79458755b06642/execroot/__main__/bazel-out/k8-opt/bin/tests/lobpcg_test_cpu.runfiles/pypi_absl_py/site-packages/absl/testing/parameterized.py", line 319, in bound_param_test
    return test_method(self, **testcase_params)
  File "/root/.cache/bazel/_bazel_root/840503f2165a538d6d79458755b06642/execroot/__main__/bazel-out/k8-opt/bin/tests/lobpcg_test_cpu.runfiles/__main__/tests/lobpcg_test.py", line 408, in testLobpcgMonotonicityF32
    self.checkLobpcgMonotonicity(matrix_name, n, k, m, tol, jnp.float32)
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/.cache/bazel/_bazel_root/840503f2165a538d6d79458755b06642/execroot/__main__/bazel-out/k8-opt/bin/tests/lobpcg_test_cpu.runfiles/__main__/tests/lobpcg_test.py", line 272, in checkLobpcgMonotonicity
    self._possibly_plot(A, eigs, X, m, matrix_name)
    ~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/.cache/bazel/_bazel_root/840503f2165a538d6d79458755b06642/execroot/__main__/bazel-out/k8-opt/bin/tests/lobpcg_test_cpu.runfiles/__main__/tests/lobpcg_test.py", line 290, in _possibly_plot
    self._debug_plots(X, eigs, info, matrix_name, plot_dir)
    ~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/.cache/bazel/_bazel_root/840503f2165a538d6d79458755b06642/execroot/__main__/bazel-out/k8-opt/bin/tests/lobpcg_test_cpu.runfiles/__main__/tests/lobpcg_test.py", line 318, in _debug_plots
    ax0.legend()
    ~~~~~~~~~~^^
  File "/root/.cache/bazel/_bazel_root/840503f2165a538d6d79458755b06642/execroot/__main__/bazel-out/k8-opt/bin/tests/lobpcg_test_cpu.runfiles/pypi_matplotlib/site-packages/matplotlib/axes/_axes.py", line 337, in legend
    self.legend_ = mlegend.Legend(self, handles, labels, **kwargs)
                   ~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/.cache/bazel/_bazel_root/840503f2165a538d6d79458755b06642/execroot/__main__/bazel-out/k8-opt/bin/tests/lobpcg_test_cpu.runfiles/pypi_matplotlib/site-packages/matplotlib/legend.py", line 549, in __init__
    self._init_legend_box(handles, labels, markerfirst)
    ~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/.cache/bazel/_bazel_root/840503f2165a538d6d79458755b06642/execroot/__main__/bazel-out/k8-opt/bin/tests/lobpcg_test_cpu.runfiles/pypi_matplotlib/site-packages/matplotlib/legend.py", line 896, in _init_legend_box
    handle_list.append(handler.legend_artist(self, orig_handle,
                       ~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^
                                             fontsize, handlebox))
                                             ^^^^^^^^^^^^^^^^^^^^
  File "/root/.cache/bazel/_bazel_root/840503f2165a538d6d79458755b06642/execroot/__main__/bazel-out/k8-opt/bin/tests/lobpcg_test_cpu.runfiles/pypi_matplotlib/site-packages/matplotlib/legend_handler.py", line 129, in legend_artist
    artists = self.create_artists(legend, orig_handle,
                                  xdescent, ydescent, width, height,
                                  fontsize, handlebox.get_transform())
  File "/root/.cache/bazel/_bazel_root/840503f2165a538d6d79458755b06642/execroot/__main__/bazel-out/k8-opt/bin/tests/lobpcg_test_cpu.runfiles/pypi_matplotlib/site-packages/matplotlib/legend_handler.py", line 303, in create_artists
    self.update_prop(legline, orig_handle, legend)
    ~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/.cache/bazel/_bazel_root/840503f2165a538d6d79458755b06642/execroot/__main__/bazel-out/k8-opt/bin/tests/lobpcg_test_cpu.runfiles/pypi_matplotlib/site-packages/matplotlib/legend_handler.py", line 88, in update_prop
    self._update_prop(legend_handle, orig_handle)
    ~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/.cache/bazel/_bazel_root/840503f2165a538d6d79458755b06642/execroot/__main__/bazel-out/k8-opt/bin/tests/lobpcg_test_cpu.runfiles/pypi_matplotlib/site-packages/matplotlib/legend_handler.py", line 79, in _update_prop
    self._default_update_prop(legend_handle, orig_handle)
    ~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/.cache/bazel/_bazel_root/840503f2165a538d6d79458755b06642/execroot/__main__/bazel-out/k8-opt/bin/tests/lobpcg_test_cpu.runfiles/pypi_matplotlib/site-packages/matplotlib/legend_handler.py", line 84, in _default_update_prop
    legend_handle.update_from(orig_handle)
    ~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^
  File "/root/.cache/bazel/_bazel_root/840503f2165a538d6d79458755b06642/execroot/__main__/bazel-out/k8-opt/bin/tests/lobpcg_test_cpu.runfiles/pypi_matplotlib/site-packages/matplotlib/lines.py", line 1358, in update_from
    self._marker = MarkerStyle(marker=other._marker)
                   ~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^
  File "/root/.cache/bazel/_bazel_root/840503f2165a538d6d79458755b06642/execroot/__main__/bazel-out/k8-opt/bin/tests/lobpcg_test_cpu.runfiles/pypi_matplotlib/site-packages/matplotlib/markers.py", line 248, in __init__
    self._set_marker(marker)
    ~~~~~~~~~~~~~~~~^^^^^^^^
  File "/root/.cache/bazel/_bazel_root/840503f2165a538d6d79458755b06642/execroot/__main__/bazel-out/k8-opt/bin/tests/lobpcg_test_cpu.runfiles/pypi_matplotlib/site-packages/matplotlib/markers.py", line 323, in _set_marker
    self.__dict__ = copy.deepcopy(marker.__dict__)
                    ~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
  File "/root/.cache/bazel/_bazel_root/840503f2165a538d6d79458755b06642/external/python_x86_64-unknown-linux-gnu-freethreaded/lib/python3.14t/copy.py", line 131, in deepcopy
    y = copier(x, memo)
  File "/root/.cache/bazel/_bazel_root/840503f2165a538d6d79458755b06642/external/python_x86_64-unknown-linux-gnu-freethreaded/lib/python3.14t/copy.py", line 202, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
                             ~~~~~~~~^^^^^^^^^^^^^
  File "/root/.cache/bazel/_bazel_root/840503f2165a538d6d79458755b06642/external/python_x86_64-unknown-linux-gnu-freethreaded/lib/python3.14t/copy.py", line 138, in deepcopy
    y = copier(memo)
  File "/root/.cache/bazel/_bazel_root/840503f2165a538d6d79458755b06642/execroot/__main__/bazel-out/k8-opt/bin/tests/lobpcg_test_cpu.runfiles/pypi_matplotlib/site-packages/matplotlib/path.py", line 285, in __deepcopy__
    p = copy.deepcopy(super(), memo)
  File "/root/.cache/bazel/_bazel_root/840503f2165a538d6d79458755b06642/external/python_x86_64-unknown-linux-gnu-freethreaded/lib/python3.14t/copy.py", line 157, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/root/.cache/bazel/_bazel_root/840503f2165a538d6d79458755b06642/external/python_x86_64-unknown-linux-gnu-freethreaded/lib/python3.14t/copy.py", line 234, in _reconstruct
    y = func(*args)
  File "/root/.cache/bazel/_bazel_root/840503f2165a538d6d79458755b06642/external/python_x86_64-unknown-linux-gnu-freethreaded/lib/python3.14t/copy.py", line 233, in <genexpr>
    args = (deepcopy(arg, memo) for arg in args)
            ~~~~~~~~^^^^^^^^^^^
  File "/root/.cache/bazel/_bazel_root/840503f2165a538d6d79458755b06642/external/python_x86_64-unknown-linux-gnu-freethreaded/lib/python3.14t/copy.py", line 138, in deepcopy
    y = copier(memo)
  File "/root/.cache/bazel/_bazel_root/840503f2165a538d6d79458755b06642/execroot/__main__/bazel-out/k8-opt/bin/tests/lobpcg_test_cpu.runfiles/pypi_matplotlib/site-packages/matplotlib/path.py", line 285, in __deepcopy__
    p = copy.deepcopy(super(), memo)
```
…pping the entire function in auto axes.

PiperOrigin-RevId: 743148311
…MIT_DEBUG_PLOTS

PiperOrigin-RevId: 743171850
Bumps [actions/cache](https://github.com/actions/cache) from 4.2.0 to 4.2.3.
- [Release notes](https://github.com/actions/cache/releases)
- [Changelog](https://github.com/actions/cache/blob/main/RELEASES.md)
- [Commits](actions/cache@1bd1e32...5a3ec84)

---
updated-dependencies:
- dependency-name: actions/cache
  dependency-type: direct:production
  update-type: version-update:semver-patch
...

Signed-off-by: dependabot[bot] <[email protected]>
dimitar-asenov and others added 26 commits April 9, 2025 06:11
…k_rep

this was essentially another instance of the jax-ml#27797 fix

fixes jax-ml#24762
…er than all-gather and print the entire logical arguments.

PiperOrigin-RevId: 745678930
…_tuple_replace

PiperOrigin-RevId: 745683630
also resolves a TODO(mattjj,frostig) from jax-ml#14570 to make vmap-of-custom_jvp not instantiate symbolic zeros

fixes jax-ml#26763
PiperOrigin-RevId: 745704582
Description:
- Added update_one_slot race tsan suppression for 3.14 TSAN CI job
  - python/cpython#132214
…ns of StableHLO.

PiperOrigin-RevId: 745717137
This function is unused.

PiperOrigin-RevId: 745721419
Co-authored-by: Yash Katariya <[email protected]>
…ne_slot-suppress

PiperOrigin-RevId: 745740772
Co-authored-by: Matthew Johnson <[email protected]>
PiperOrigin-RevId: 745796269
Move methods out of line, separate public and private sections.

Readability cleanup only, no functional changes intended.

PiperOrigin-RevId: 745802620
Refactoring only, no functional changes intended.

PiperOrigin-RevId: 745813442
@rocm-repo-management-api-2 rocm-repo-management-api-2 bot requested a review from a team as a code owner April 10, 2025 06:02
@rocm-repo-management-api-2 rocm-repo-management-api-2 bot enabled auto-merge (rebase) April 10, 2025 06:02
auto-merge was automatically disabled May 1, 2025 15:11

Pull request was closed

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.