forked from jax-ml/jax
-
Notifications
You must be signed in to change notification settings - Fork 4
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
Closed
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
PiperOrigin-RevId: 742840628
…allas_2 PiperOrigin-RevId: 742860659
PiperOrigin-RevId: 742864570
PiperOrigin-RevId: 742878778
PiperOrigin-RevId: 742886445
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: 743046047
PiperOrigin-RevId: 743054832
PiperOrigin-RevId: 743054850
…-primitive PiperOrigin-RevId: 743059662
Reverts 33d306a PiperOrigin-RevId: 743062185
Reverts 0b199f4 PiperOrigin-RevId: 743063615
PiperOrigin-RevId: 743074472
http://github.com/openxla/xla/commit/c3087e022f3c07f7ed1dd4e47024c437a504341b. PiperOrigin-RevId: 743093178
…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
PiperOrigin-RevId: 743171768
…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]>
PiperOrigin-RevId: 743172056
PiperOrigin-RevId: 745557539
http://github.com/openxla/xla/commit/f2a101a0a20c0ebadc15e442f2e0923fc8e78e8f. PiperOrigin-RevId: 745562488
PiperOrigin-RevId: 745621495
…k_rep this was essentially another instance of the jax-ml#27797 fix fixes jax-ml#24762
PiperOrigin-RevId: 745647632
…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: 745692664
PiperOrigin-RevId: 745704582
PiperOrigin-RevId: 745707421
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]>
PiperOrigin-RevId: 745732217
PiperOrigin-RevId: 745739477
…ne_slot-suppress PiperOrigin-RevId: 745740772
PiperOrigin-RevId: 745742774
Co-authored-by: Matthew Johnson <[email protected]> PiperOrigin-RevId: 745796269
PiperOrigin-RevId: 745799649
Move methods out of line, separate public and private sections. Readability cleanup only, no functional changes intended. PiperOrigin-RevId: 745802620
PiperOrigin-RevId: 745813055
Refactoring only, no functional changes intended. PiperOrigin-RevId: 745813442
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Daily sync with upstream