ROCm/jax

Pull Requests Commits

Update XLA dependency to use revision

Google-ML-Automation committed 340 days ago

7a459f0e

[Mosaic GPU] Add transform inference rule for `mgpu.slice_smem`.

bchetioui committed 340 days ago

ba2f7c9a

[Mosaic GPU] Add a specialized layout for loading 4-bit inputs in WGMMA friendly layouts

apaszke committed 340 days ago

d4bd2570

[mosaic_gpu] Force flush all cupti activity, then unsubscribe.

chr1sj0nes committed 340 days ago

38d52a19

[Mosaic GPU] Remove sub-byte conversion restriction

apaszke committed 340 days ago

34cd5b0d

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

yashk2810 committed 340 days ago

549973de

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

changhuilin committed 340 days ago

f174b00f

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

emilyfertig committed 340 days ago

8c351917

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

superbobry committed 340 days ago

051687dc

Compute tile index using tile-based coordinates

Google-ML-Automation committed 340 days ago

b4966130

Merge pull request #27164 from MichaelHudgins:a4-testing

Google-ML-Automation committed 340 days ago

b74b16f9

Add B200 testing to continuous workflow

MichaelHudgins committed 340 days ago

ecf7fde7

Fix error in pallas tutorial

Google-ML-Automation committed 340 days ago

4f704713

Replace cached function get_replicated_hlo_sharding() with a constant.

hawkinsp committed 340 days ago

20658fab

Merge pull request #26980 from carlosgmartin:categorical_replace

Google-ML-Automation committed 340 days ago

ebcae0d3

Remove code that preserved _original_py_fns on C++ classes.

hawkinsp committed 340 days ago

be5d13af

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

bchetioui committed 340 days ago

9a686e0b

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

carlosgmartin committed 340 days ago

3f59fa68

Fix auditwheel version issue (#288)

mrodden committed 341 days ago

Verified d864b4fb

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

Google-ML-Automation committed 341 days ago

de9ad6ba

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

nitins17 committed 341 days ago

031614c2

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

apaszke committed 341 days ago

3649da56

Removed trivial docstrings from JAX tests

superbobry committed 341 days ago

0ff23404

Update XLA dependency to use revision

Google-ML-Automation committed 341 days ago

55812c5d

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

superbobry committed 341 days ago

a7e5eaee

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

apaszke committed 341 days ago

89b21de6

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

apaszke committed 341 days ago

2bdd9c87

Merge pull request #27176 from jakevdp:lax-docs

Google-ML-Automation committed 342 days ago

761b35c5

Update XLA dependency to use revision

Google-ML-Automation committed 342 days ago

e8b683ae

Change the way that batching.spec_types is updated.

jpuigcerver committed 342 days ago

466ef6a1

Newer Older