CI: 03/17/25 upstream sync #292
charleshofer
merged
23 commits into
rocm-main
from
ci-upstream-sync-149_1
Use `lax.top_k` instead of `jnp.argsort` in Gumbel top-k trick for we…
4a82fe94
add experimental lax.optimization_barrier autodiff rules
dadc68b6
Add a C++ implementation of a toplogical sort.
14cb7453
Merge pull request #27174 from mattjj:opt-barrier-ad-rules
7db59cdc
mixing modes
3c0027af
Merge pull request #27177 from jax-ml:mixing_modes
d07d642d
Support error checking in explicit mode
9b0ace4a
Update XLA dependency to use revision
f360e191
Better docs for jax.lax add/sub/mul/div
de8b0564
Change the way that batching.spec_types is updated.
466ef6a1
Update XLA dependency to use revision
e8b683ae
Merge pull request #27176 from jakevdp:lax-docs
761b35c5
[Mosaic GPU] Add support for fast WGMMA layout changes after 8- to 16…
2bdd9c87
[Mosaic GPU] Add support for changing the layout before the upcast
89b21de6
[pallas:mosaic_gpu] `jnp.reduce_sum` now works for >1D arrays
a7e5eaee
Update XLA dependency to use revision
55812c5d
Removed trivial docstrings from JAX tests
0ff23404
[Mosaic GPU] Make the s4 -> bf16 upcast more flexible when it comes t…
3649da56
Pin numpy~=2.1.0 in workflow file instead of test-requirements.txt
031614c2
Merge pull request #27157 from mar-muel:improve-random-choice-perform…
de9ad6ba
[Mosaic GPU] Add initial transform inference rules for `vector.{load,…
9a686e0b
Merge remote-tracking branch 'origin/rocm-main' into ci-upstream-sync…
0c3b9814
Trivial change for CI
d1b79232
disabled auto-merge 289 days ago
Manually disabled by user
Manually disabled by user
charleshofer
approved these changes
on 2025-03-17
Login to write a write a comment.
Login via GitHub