ROCm/jax

Pull Requests Commits

Enable upstream CI on release branches

charleshofer committed 330 days ago

c4f188cc

Merge pull request #294 from ROCm/ci-upstream-sync-151_1

charleshofer committed 330 days ago

Verified c46b4fc0

Merge branch 'rocm-main' into ci-upstream-sync-151_1

charleshofer committed 330 days ago

c7b407c9

[Mosaic GPU] Raise a `NotImplementedError` if `swizzle=16`.

bchetioui committed 330 days ago

1e36cbe5

[Mosaic GPU] Fuse slicing into s4 -> bf16 upcasts

apaszke committed 330 days ago

8da93249

Update XLA dependency to use revision

Google-ML-Automation committed 330 days ago

7a459f0e

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

bchetioui committed 330 days ago

ba2f7c9a

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

apaszke committed 330 days ago

d4bd2570

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

chr1sj0nes committed 330 days ago

38d52a19

[Mosaic GPU] Remove sub-byte conversion restriction

apaszke committed 330 days ago

34cd5b0d

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

yashk2810 committed 331 days ago

549973de

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

changhuilin committed 331 days ago

f174b00f

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

emilyfertig committed 331 days ago

8c351917

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

superbobry committed 331 days ago

051687dc

Compute tile index using tile-based coordinates

Google-ML-Automation committed 331 days ago

b4966130

Merge pull request #27164 from MichaelHudgins:a4-testing

Google-ML-Automation committed 331 days ago

b74b16f9

Add B200 testing to continuous workflow

MichaelHudgins committed 331 days ago

ecf7fde7

Fix error in pallas tutorial

Google-ML-Automation committed 331 days ago

4f704713

Replace cached function get_replicated_hlo_sharding() with a constant.

hawkinsp committed 331 days ago

20658fab

Merge pull request #26980 from carlosgmartin:categorical_replace

Google-ML-Automation committed 331 days ago

ebcae0d3

Remove code that preserved _original_py_fns on C++ classes.

hawkinsp committed 331 days ago

be5d13af

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

bchetioui committed 331 days ago

9a686e0b

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

carlosgmartin committed 331 days ago

3f59fa68

Fix auditwheel version issue (#288)

mrodden committed 331 days ago

Verified d864b4fb

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

Google-ML-Automation committed 331 days ago

de9ad6ba

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

nitins17 committed 331 days ago

031614c2

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

apaszke committed 331 days ago

3649da56

Removed trivial docstrings from JAX tests

superbobry committed 331 days ago

0ff23404

Update XLA dependency to use revision

Google-ML-Automation committed 331 days ago

55812c5d

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

superbobry committed 331 days ago

a7e5eaee

Older