forked from jax-ml/jax
-
Notifications
You must be signed in to change notification settings - Fork 4
CI: 04/02/25 upstream sync #335
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: 739874654
Very importantly, this also includes support for loading the packed accumulator from TMEM. PiperOrigin-RevId: 739883035
http://github.com/openxla/xla/commit/9a8dd0796bcfeb00e4e6d09d74726db5c7d4a003. PiperOrigin-RevId: 739886693
`TiledLayout` has no notion of partitioning over warpgroups, and each warp holds `16 x 8` elements. PiperOrigin-RevId: 739942481
PiperOrigin-RevId: 739946581
PiperOrigin-RevId: 739952849
…version. PiperOrigin-RevId: 739958612
PiperOrigin-RevId: 739967341
… Triton lowering. PiperOrigin-RevId: 739968019
PiperOrigin-RevId: 739968243
PiperOrigin-RevId: 739969862
…ed layout. PiperOrigin-RevId: 739995908
Reverts 53e8eac PiperOrigin-RevId: 740012608
These args are optional, so not specifying them in our tests will make them simpler and easier to read. This change is a no-op. PiperOrigin-RevId: 740015584
The xla_extension move introduced an incorrect path.
…onv)`. This change is a no-op made for convenience for follow-up changes. PiperOrigin-RevId: 740020625
Kidger's reimplementation of this code notes that the break mode and indent are unused in the _fits function (https://github.com/patrick-kidger/wadler_lindig/blob/851379b8f55e2bb98ea2c81905863f90f9606f0a/wadler_lindig/_wadler_lindig.py#L166). We can make the same optimization here.
PiperOrigin-RevId: 740033782
Remove jax dependencies from the Bazel test targets for `:build_jaxlib=false` and `:build_jaxlib=wheel`. `internal_test_util` is removed from the `jax` wheel. To use this package in Bazel py_test, we need to copy it to the unpacked wheel folder. This is done by adding `wheel_deps` value to `py_import` Jax targets. This change concludes ML Wheels design implementation in JAX and enables testing of all wheels via Bazel command. PiperOrigin-RevId: 740037952
PiperOrigin-RevId: 740041921
PiperOrigin-RevId: 740043785
PiperOrigin-RevId: 740050893
PiperOrigin-RevId: 740051013
…d a 1d array and broadcast it along the leading dimension to a 2d shape as an input to a wgmma. In this new layout the first 4 threads of a warp group hold 8 uniques values. These values are replicated in each (thread_idx % 4) group. PiperOrigin-RevId: 740058172
Indexing is less verbose and is thus easier to read in most cases. The functional API is really only necessary for masked loads and stores. PiperOrigin-RevId: 740058341
…ension and load a ref into WGMMAColFragLayout format. PiperOrigin-RevId: 740068368
PiperOrigin-RevId: 740084295
PiperOrigin-RevId: 742447929
PiperOrigin-RevId: 742464022
…nversions This will be submitted automatically once the compatibility window has passed PiperOrigin-RevId: 742464046
PiperOrigin-RevId: 742531956
http://github.com/openxla/xla/commit/b1971cc2b3407e87fada2674a057d72897b79acc. PiperOrigin-RevId: 742646393
Add the tests to the Bazel presubmit RBE jobs (except `arm64`/`aarch64` jobs that use RBE cross-compilation). PiperOrigin-RevId: 742724458
PiperOrigin-RevId: 742737384
…/tsickert/discord-webhook-7.0.0 PiperOrigin-RevId: 742756295
PiperOrigin-RevId: 742771594
PiperOrigin-RevId: 742774171
Shardy has been been included in JAX for a while now. PiperOrigin-RevId: 742778405
Reverts 12526ea PiperOrigin-RevId: 742784183
PiperOrigin-RevId: 742795353
PiperOrigin-RevId: 742808140
PiperOrigin-RevId: 742840628
…allas_2 PiperOrigin-RevId: 742860659
PiperOrigin-RevId: 742864570
PiperOrigin-RevId: 742878778
PiperOrigin-RevId: 742886445
PiperOrigin-RevId: 742927160
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