ROCm/jax

Pull Requests Commits

Trivial change for CI

charleshofer committed 328 days ago

d1b79232

Merge remote-tracking branch 'origin/rocm-main' into ci-upstream-sync-149_1

invalid-email-address committed 328 days ago

0c3b9814

[Mosaic GPU] Add initial transform inference rules for `vector.{load,store}`.

bchetioui committed 328 days ago

9a686e0b

Fix auditwheel version issue (#288)

mrodden committed 328 days ago

Verified d864b4fb

Merge pull request #27157 from mar-muel:improve-random-choice-performance

Google-ML-Automation committed 328 days ago

de9ad6ba

Pin numpy~=2.1.0 in workflow file instead of test-requirements.txt

nitins17 committed 328 days ago

031614c2

[Mosaic GPU] Make the s4 -> bf16 upcast more flexible when it comes to vector length

apaszke committed 328 days ago

3649da56

Removed trivial docstrings from JAX tests

superbobry committed 328 days ago

0ff23404

Update XLA dependency to use revision

Google-ML-Automation committed 328 days ago

55812c5d

[pallas:mosaic_gpu] `jnp.reduce_sum` now works for >1D arrays

superbobry committed 328 days ago

a7e5eaee

[Mosaic GPU] Add support for changing the layout before the upcast

apaszke committed 329 days ago

89b21de6

[Mosaic GPU] Add support for fast WGMMA layout changes after 8- to 16-bit upcast

apaszke committed 329 days ago

2bdd9c87

Merge pull request #27176 from jakevdp:lax-docs

Google-ML-Automation committed 329 days ago

761b35c5

Update XLA dependency to use revision

Google-ML-Automation committed 329 days ago

e8b683ae

Change the way that batching.spec_types is updated.

jpuigcerver committed 330 days ago

466ef6a1

Better docs for jax.lax add/sub/mul/div

jakevdp committed 330 days ago

de8b0564

Update XLA dependency to use revision

Google-ML-Automation committed 330 days ago

f360e191

Support error checking in explicit mode

ayaka14732 committed 331 days ago

9b0ace4a

Merge pull request #27177 from jax-ml:mixing_modes

Google-ML-Automation committed 331 days ago

d07d642d

yashk2810 committed 331 days ago

3c0027af

Merge pull request #27174 from mattjj:opt-barrier-ad-rules

Google-ML-Automation committed 331 days ago

7db59cdc

Add a C++ implementation of a toplogical sort.

hawkinsp committed 331 days ago

14cb7453

Merge pull request #285 from ROCm/ci-upstream-sync-147_1

github-actions[bot] committed 331 days ago

Verified 7a172d70

Trivial change for CI

charleshofer committed 331 days ago

7a6940bb

Merge remote-tracking branch 'origin/rocm-main' into ci-upstream-sync-147_1

invalid-email-address committed 331 days ago

e275d5cf

add experimental lax.optimization_barrier autodiff rules

mattjj committed 331 days ago

dadc68b6

Merge pull request #27015 from mattjj:direct-linearize-fixes-4

Google-ML-Automation committed 331 days ago

b00a3a19

[pallas:mosaic_gpu] WG lowering now supports `while_p`

superbobry committed 331 days ago

64230d1c

Count test totals correctly for dashboards (#280)

charleshofer committed 331 days ago

Verified 022da913

[direct-linearize] shmap fixes

mattjj committed 331 days ago

174dcc77

Older