ROCm/jax

Pull Requests Commits

Trivial change for CI

charleshofer committed 1 year ago

d1b79232

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

invalid-email-address committed 1 year ago

0c3b9814

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

bchetioui committed 1 year ago

9a686e0b

Fix auditwheel version issue (#288)

mrodden committed 1 year ago

Verified d864b4fb

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

Google-ML-Automation committed 1 year ago

de9ad6ba

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

nitins17 committed 1 year ago

031614c2

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

apaszke committed 1 year ago

3649da56

Removed trivial docstrings from JAX tests

superbobry committed 1 year ago

0ff23404

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

55812c5d

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

superbobry committed 1 year ago

a7e5eaee

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

apaszke committed 1 year ago

89b21de6

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

apaszke committed 1 year ago

2bdd9c87

Merge pull request #27176 from jakevdp:lax-docs

Google-ML-Automation committed 1 year ago

761b35c5

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

e8b683ae

Change the way that batching.spec_types is updated.

jpuigcerver committed 1 year ago

466ef6a1

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

jakevdp committed 1 year ago

de8b0564

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

f360e191

Support error checking in explicit mode

ayaka14732 committed 1 year ago

9b0ace4a

Merge pull request #27177 from jax-ml:mixing_modes

Google-ML-Automation committed 1 year ago

d07d642d

yashk2810 committed 1 year ago

3c0027af

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

Google-ML-Automation committed 1 year ago

7db59cdc

Add a C++ implementation of a toplogical sort.

hawkinsp committed 1 year ago

14cb7453

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

github-actions[bot] committed 1 year ago

Verified 7a172d70

Trivial change for CI

charleshofer committed 1 year ago

7a6940bb

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

invalid-email-address committed 1 year ago

e275d5cf

add experimental lax.optimization_barrier autodiff rules

mattjj committed 1 year ago

dadc68b6

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

Google-ML-Automation committed 1 year ago

b00a3a19

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

superbobry committed 1 year ago

64230d1c

Count test totals correctly for dashboards (#280)

charleshofer committed 1 year ago

Verified 022da913

[direct-linearize] shmap fixes

mattjj committed 1 year ago

174dcc77

Older