CI: 04/22/25 upstream sync #377
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}
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: 746098316
PiperOrigin-RevId: 746100873
PiperOrigin-RevId: 746102268
PiperOrigin-RevId: 746112248
PiperOrigin-RevId: 746117643
PiperOrigin-RevId: 746128775
PiperOrigin-RevId: 746140286
This fixes some non-intuitive errors where scalar-shaped values in VREGs were being used in operations that expected SREGs. PiperOrigin-RevId: 746146037
…pes being enabled by default PiperOrigin-RevId: 746146834
…sult jax.Array. PiperOrigin-RevId: 746147571
Adds a new WarpMesh object which when used in conjunction with core_map, allows the user to drop into warp-level code rather than programming at the warpgroup level. PiperOrigin-RevId: 746163942
This change primarily reduces sharding, although in a few cases it also increases shardings. It is harmful to performance to overshard tests since there's a startup and teardown cost to each test run. In a few cases, change tests to be non-accelerator tests. PiperOrigin-RevId: 746164539
Partially addresses: jax-ml#18246. If compile can also be a future, this code can be used to safely block on that as well. PiperOrigin-RevId: 746189742
…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
PiperOrigin-RevId: 749779206
http://github.com/openxla/xla/commit/9219fd7ef180a01f814d3fce9f8aecfd80b9fd6c. PiperOrigin-RevId: 749784566
Description: - Copy mlir module before adding new attributes Fixes jax-ml#27991
PiperOrigin-RevId: 749807401
…utation-27991 PiperOrigin-RevId: 749811476
PiperOrigin-RevId: 749813377
…ed op instead of multiple .at[] calls. PiperOrigin-RevId: 749818535
Amend the scheme format and top-level domain.
PiperOrigin-RevId: 749885945
PiperOrigin-RevId: 749889362
PiperOrigin-RevId: 749898586
PiperOrigin-RevId: 749908416
Migrate auto-tuned table from https://github.com/pytorch/xla/blob/master/torch_xla/experimental/tuned_block_sizes.py PiperOrigin-RevId: 749965181
…bstract eval This can happen if a user forgets to unwrap a ref! @asabne had this happen to him today, and he was confused as to what was going on. The prior error is unclear: AssertionError: (MemRef<None>{float32[2,1024,1024]}, MemRef<None>{float32[1,1024,1024]}) PiperOrigin-RevId: 749979253
…ray creation when possible This changes makes use of the new `xla::ifrt::Client::MakeArraysFromHostBufferShards()` API when possible. This API needs a single call to create a multi-shard IFRT Array (to be wrapped as a JAX `PyArray`), which provides more optimization opportunities for the runtime than creating single-device IFRT Arrays and then assembling them. Please note that `xla::ifrt::Client::MakeArraysFromHostBufferShards()` implementation in PjRt-IFRT is not yet optimized, so there is no immediate performance benefits for McJAX. As an exception, it takes the previous path of array assembly if any shard for `BatchedDevicePut` is not a host buffer, but already a single-device array, because `xla::ifrt::Client::MakeArraysFromHostBufferShards()` works only if all the sharded input to be host buffers. With batching possible at IFRT level, we now skip `DevicePutResultFn` step; `DevicePut` (now `DevicePutWithDevice` and `DevicePutWithSharding`) internally calls per-shard functions (with GIL released) and returns a final IFRT Array. This change includes a code cleanup for `xla::DevicePutResult::owning_pybuffer`, which was originally intended to hold a Python object to keep an IFRT Array valid when it is created from `DevicePut()` implementations, but this role has been entirely covered by `on_done_with_host_buffer` function supplied at IFRT Array creation time. PiperOrigin-RevId: 749989229
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