ROCm/jax

Pull Requests Commits

Allow pspec to be passed to device_put if there is a mesh in the surrounding context

yashk2810 committed 343 days ago

549973de

Replace the uses of `PjRtClient::Compile()` with `PjRtClient::CompileAndLoad()`.

changhuilin committed 343 days ago

f174b00f

Enable `jax.device_put` to a sharding with no local devices.

emilyfertig committed 343 days ago

8c351917

[pallas] `pallas_call_p` is now parameterized by a mesh

superbobry committed 343 days ago

051687dc

Compute tile index using tile-based coordinates

Google-ML-Automation committed 343 days ago

b4966130

Merge pull request #27164 from MichaelHudgins:a4-testing

Google-ML-Automation committed 343 days ago

b74b16f9

Add B200 testing to continuous workflow

MichaelHudgins committed 343 days ago

ecf7fde7

Fix error in pallas tutorial

Google-ML-Automation committed 343 days ago

4f704713

Replace cached function get_replicated_hlo_sharding() with a constant.

hawkinsp committed 343 days ago

20658fab

Merge pull request #26980 from carlosgmartin:categorical_replace

Google-ML-Automation committed 343 days ago

ebcae0d3

Remove code that preserved _original_py_fns on C++ classes.

hawkinsp committed 343 days ago

be5d13af

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

bchetioui committed 343 days ago

9a686e0b

Add replace option to random.categorical to enable sampling without replacement.

carlosgmartin committed 343 days ago

3f59fa68

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

Google-ML-Automation committed 343 days ago

de9ad6ba

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

nitins17 committed 343 days ago

031614c2

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

apaszke committed 343 days ago

3649da56

Removed trivial docstrings from JAX tests

superbobry committed 343 days ago

0ff23404

Update XLA dependency to use revision

Google-ML-Automation committed 343 days ago

55812c5d

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

superbobry committed 343 days ago

a7e5eaee

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

apaszke committed 343 days ago

89b21de6

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

apaszke committed 343 days ago

2bdd9c87

Merge pull request #27176 from jakevdp:lax-docs

Google-ML-Automation committed 344 days ago

761b35c5

Update XLA dependency to use revision

Google-ML-Automation committed 344 days ago

e8b683ae

Change the way that batching.spec_types is updated.

jpuigcerver committed 344 days ago

466ef6a1

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

jakevdp committed 345 days ago

de8b0564

Update XLA dependency to use revision

Google-ML-Automation committed 345 days ago

f360e191

Support error checking in explicit mode

ayaka14732 committed 346 days ago

9b0ace4a

Merge pull request #27177 from jax-ml:mixing_modes

Google-ML-Automation committed 346 days ago

d07d642d

yashk2810 committed 346 days ago

3c0027af

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

Google-ML-Automation committed 346 days ago

7db59cdc

Newer Older