CI: 04/11/25 upstream sync #356
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
![rocm-repo-management-api-2[bot]](https://avatars.githubusercontent.com/u/21157610?s=60&v=4){kind=small}
PiperOrigin-RevId: 743342672
PiperOrigin-RevId: 743343131
Drop into `Auto` mode in the implementation.
PiperOrigin-RevId: 743441718
PiperOrigin-RevId: 743452172
PiperOrigin-RevId: 743452501
PiperOrigin-RevId: 743478967
…ly added primitives PiperOrigin-RevId: 743492343
… the body This replaces the old behavior where `emit_pipeline*` would replace the current parallel grid with the sequential grid, changing the output of `pl.program_id`. With this change, `pl.program_id` always works wrt the parallel grid. PiperOrigin-RevId: 743498194
This is no longer used, since we elected to refine the IR by annotating it with
`{in,out}_transforms` in the lowering pipeline instead.
PiperOrigin-RevId: 743516621
http://github.com/openxla/xla/commit/921c164a67e8ac4cf052aab26e849f29b719f802. PiperOrigin-RevId: 743535272
It's important we use the un-stemmed name because on Windows there is an .exe suffix.
PiperOrigin-RevId: 743549467
The change is made to address the case when bazel dir has multiple wheels with different version suffixes. We need to copy only those wheels that were created by the current execution of build.py script. PiperOrigin-RevId: 743566122
This parameter is already part of the internal API for the AOT lowering function, here we just expose it to `jax.export`.
This allows `ArrayImpl` to directly subclass `jax.Array` without relying on the expensive virtual subclasses machinery from `abc`. PiperOrigin-RevId: 743573028
There's a mismatch between the resulting and the desired matrixes on H100, but not the older GPUs. PiperOrigin-RevId: 743578025
We found a hanging issue when we use dynamic grid. We'll disable it for now. PiperOrigin-RevId: 743597352
PiperOrigin-RevId: 743601775
PiperOrigin-RevId: 743611388
PiperOrigin-RevId: 743611394
PiperOrigin-RevId: 743633331
PiperOrigin-RevId: 743658654
The code was using ~ with a boolean, which leads to a new DeprecationWarning. That should only be used with ints. PiperOrigin-RevId: 743668386
PiperOrigin-RevId: 743675002
PiperOrigin-RevId: 743727350
Replace `_handle_reshape()` and `_handle_index()` with a general
`_handle_transform()` that applies all transforms except tiling and (optionally)
transposes. The implementation is based on
`_untransform_{transpose,reshape,index}()` transform methods on transforms that
find the conjugate of the transpose/reshape/index wrt the transform.
PiperOrigin-RevId: 743731515
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
Also remove some unnecessary LINKOPTS. These are no longer needed now we use the pywrap rules instead. PiperOrigin-RevId: 746216832
PiperOrigin-RevId: 746217405
We are seeing a higher number of cancellations of the continuous job recently: ``` Canceling since a higher priority waiting request for CI - Wheel Tests (Continuous)-refs/heads/main exists ``` PiperOrigin-RevId: 746222323
PiperOrigin-RevId: 746226562
PiperOrigin-RevId: 746228503
also fix a bug in ordered effects in cond_fun lowering fixes google/flax#4700
…mization PiperOrigin-RevId: 746250385
We need v1.15.2 for Linux aarch64 3.13-t wheels to exist.
PiperOrigin-RevId: 746257919
Fixes jax-ml#22783 jax-fixit PiperOrigin-RevId: 746260869
PiperOrigin-RevId: 746295338
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