CI: 04/21/25 upstream sync #376
Open
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
![rocm-repo-management-api-2[bot]](https://avatars.githubusercontent.com/u/21157610?s=60&v=4){kind=small}
…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
Co-authored-by: Yash Katariya <[email protected]>
PiperOrigin-RevId: 745813055
Refactoring only, no functional changes intended. PiperOrigin-RevId: 745813442
…est case. This is very helpful for debugging deadlocks! PiperOrigin-RevId: 745986596
PiperOrigin-RevId: 745992513
improve layout, information add dummy import to hopefully fix build issue parse help text for markdown whoops didn't mean to do it twice jax prefix text no longer applies here two space indents address definition list ending without blank line error provide deprecation mechanism document context managagers if they exist remove mention of context manager try and fix formatting improve formatting, fail to fix warnings fail to fix bug, make better looking anyway okay bug was in the parsing of help text to rst, some of which does not parse wow, found the bug, turns out help strings were not valid rst
http://github.com/openxla/xla/commit/9f2aa85b909a615eef2212286f1e0c5684fc6b6c. PiperOrigin-RevId: 746004084
Without this, `emit_pipeline_warp_specialized` would leave the barriers in a bad state, causing deadlocks or crashes when it was called multiple times in sequence. PiperOrigin-RevId: 746022784
PiperOrigin-RevId: 746026464
PiperOrigin-RevId: 746033054
…teps in the memory WG max_concurrent_steps is an upper bound: we no longer guarantee that it accurately reflects the actual number of steps when the grid has dynamic bounds PiperOrigin-RevId: 746036125
PiperOrigin-RevId: 746056606
PiperOrigin-RevId: 746057793
If mesh axes are empty, we are setting mesh as None, resulting in an error in this test. This fix provides an empty mesh, when no mesh axes in dumped module are empty. PiperOrigin-RevId: 746058506
PiperOrigin-RevId: 746059204
PiperOrigin-RevId: 746059522
…rying with the correct vma as the operands were. PiperOrigin-RevId: 746065965
PiperOrigin-RevId: 746070501
PiperOrigin-RevId: 746097962
PiperOrigin-RevId: 748818185
PiperOrigin-RevId: 748839843
New minimum jaxlib_extension_version is 330. PiperOrigin-RevId: 748853497
…l ops inside `_gather` are in auto mode. Fix select's batching rule where `explicit_mesh_axis` that we capture in `axis_data` was not propagated properly to the `broadcast` happening in `bdim_at_front`. PiperOrigin-RevId: 748867490
PiperOrigin-RevId: 748880870
…each dim in the block shape Currently block_shape is tuple[int | None, …]. We propose generalizing block_shape to take in more types in the tuple to more generally support: * Squeeze dimension (currently None, could be pl.Squeezed()) * Unblocked: currently the entire index_map needs to be Unblocked or not. This will allow individual indices to be Blocked/Unblocked, e.g. pl.BlockSpec((pl.Unblocked(...), 512), …) * Ragged sizes: the index_map will return a pl.ds with a dynamic size (bounded by some something). For example: pl.BlockSpec((pl.DynamicSizedSlice(512), 1024), lambda i, j: (pl.ds(...), j). This will make BlockSpecs a lot more flexible and will enable things like doing arbitrary slicing in things like pipeline emitter. PiperOrigin-RevId: 748881960
PiperOrigin-RevId: 748914721
http://github.com/openxla/xla/commit/9477be9b2b3afdc2c93b7c272d00f8034750f41c. PiperOrigin-RevId: 749017950
…ape,XlaComputation}. This should not be a breaking change; the underlying definitions were already removed. PiperOrigin-RevId: 749037564
PiperOrigin-RevId: 749055960
(We leave the name exported by JAX to avoid breaking users, but fixed to its last known value.) PiperOrigin-RevId: 749070199
PiperOrigin-RevId: 749086392
…n `out_sharding` argument in `set`, use the input array's `sharding`. Since this is an update, after `.set`, the input array's sharding should be preserved. Fixes jax-ml#28111 PiperOrigin-RevId: 749089846
Fixes some test flakiness with JAX_TEST_NUM_THREADS > 1. PiperOrigin-RevId: 749089988
PiperOrigin-RevId: 749105590
… in a shard_map. This change reverts cl/731724837 ([github link](jax-ml@4997e45)), which is a temporary solution to solve the inconsistent padding on the boundary of manual computation. Now that we have a better solution cl/746600070 ([github link](openxla/xla#25080)), we revert this temporary solution. We still keep the added `test_partially_sharded_dim_with_auto` to verify the correctness. Reverts 4997e45 PiperOrigin-RevId: 749123936
…rly by adding the explicit mesh axis on dim 0 PiperOrigin-RevId: 749125322
PiperOrigin-RevId: 749159983
Array serialization in array_serialization.py contains a mixture of JAX specific serialization logic and tensorstore driver. This change separates JAX and tensorstore methods (a) making serialization more modular and (b) potentially allowing for alternative array serialization backends in the future. Additional clean-up changes include: - making ocdbt kvstore driver default in tensorstore - robustified array serialization tests especially on multi-host - explicit tensorstore array chunking to ensure chunk file size does not blow up PiperOrigin-RevId: 749175295
PiperOrigin-RevId: 749195753
…t to tracing cache after sharding_in_types config was turned on which lead to `sharding` always being available on `ShapedArray` PiperOrigin-RevId: 749206500
PiperOrigin-RevId: 749243140
http://github.com/openxla/xla/commit/1dd84dc2e7f87d79ba9f77b9874ff4a50227ad5e. PiperOrigin-RevId: 749332712
PiperOrigin-RevId: 749464614
PiperOrigin-RevId: 749492881
http://github.com/openxla/xla/commit/f94b36b783b9d955ec8cc966fc0b76cf9e265382. PiperOrigin-RevId: 749548364
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