Skip to content

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
wants to merge 369 commits into from

Conversation

rocm-repo-management-api-2[bot]
Copy link

Daily sync with upstream

chr1sj0nes and others added 30 commits March 24, 2025 03:28
PiperOrigin-RevId: 739874654
Very importantly, this also includes support for loading the
packed accumulator from TMEM.

PiperOrigin-RevId: 739883035
`TiledLayout` has no notion of partitioning over warpgroups, and each warp
holds `16 x 8` elements.

PiperOrigin-RevId: 739942481
… Triton lowering.

PiperOrigin-RevId: 739968019
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.
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: 740043785
PiperOrigin-RevId: 740050893
…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
Google-ML-Automation and others added 26 commits March 31, 2025 16:42
…nversions

This will be submitted automatically once the compatibility window has passed

PiperOrigin-RevId: 742464046
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
Shardy has been been included in JAX for a while now.

PiperOrigin-RevId: 742778405
PiperOrigin-RevId: 742840628
PiperOrigin-RevId: 742927160
@rocm-repo-management-api-2 rocm-repo-management-api-2 bot requested a review from a team as a code owner April 2, 2025 06:02
@rocm-repo-management-api-2 rocm-repo-management-api-2 bot enabled auto-merge (rebase) April 2, 2025 06:02
auto-merge was automatically disabled May 1, 2025 15:11

Pull request was closed

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.