ROCm/jax

Pull Requests Commits

Fix pallas fuser usage of call_hi_primitive_p

yashk2810 committed 100 days ago

1c781867

Remove dead code

yashk2810 committed 100 days ago

f2f1d3ab

Fix some jax tests when custom_vjp3 defaults to True. There are still 6 test targets failing.

yashk2810 committed 100 days ago

1371c2ba

Remove primitive_transpose registrations for primitives that have a transpose_fancy rule defined

yashk2810 committed 100 days ago

6481bb28

Update rules_cc version to 0.2.9

akuegel committed 100 days ago

8b3bc0b4

[Mosaic:GPU] Use 2 GPUs for collective metadata test

PatriosTheGreat committed 101 days ago

1ff638f9

[Mosaic GPU] Add support for tcgen05 MMAs that are both sparse and block-scaled

apaszke committed 101 days ago

4553a01e

Re-enable CI for B200 MGPU test + increase shard_count

apaszke committed 101 days ago

cf870cbe

[pallas:mosaic] Renamed the `KernelType` enum to `CoreType`

superbobry committed 101 days ago

1d0504b1

[Pallas:MGPU] Expose async SMEM->TMEM copies for arbitrary data

apaszke committed 101 days ago

757c92db

Update XLA dependency to use revision http://github.com/openxla/xla/commit/35a0f6468a0c2723f575fa049d92bc81c99b5451

Google-ML-Automation committed 101 days ago

4a28ce12

[Pallas] DCE all jaxprs before lowering

sharadmv committed 101 days ago

8fdc565d

Support atan2 in Pallas/Mosaic.

yarongmu-google committed 101 days ago

3d0ed0f8

[Rollback] remove jax_experimental_unsafe_xla_runtime_errors

mattjj committed 101 days ago

bb04e7a6

update_vma on AbstractRef updates its inner aval

Google-ML-Automation committed 101 days ago

11a3faf2

Pass overridden XLA repository location to CUDA non-RBE tests.

ybaturina committed 101 days ago

575ce86e

[Mosaic GPU] Remove legacy custom call API.

allanrenucci committed 101 days ago

5f99745e

[mosaic] Added `GetFuncWithCoreType`

superbobry committed 101 days ago

5c11726b

[Mosaic] Support erf.

WindQAQ committed 101 days ago

6e45cca4

Update rules_python version to 1.6.3 and gRPC to 1.78.0 in MODULE.bazel.

ybaturina committed 101 days ago

fe997775

Merge pull request #35236 from superbobry:main

Google-ML-Automation committed 101 days ago

1da5af44

[checkify] remove jax_experimental_unsafe_xla_runtime_errors

mattjj committed 101 days ago

bf610473

Merge pull request #34943 from jurasic-pf:jurasic/faster-trapezoid

Google-ML-Automation committed 101 days ago

d09caeae

Merge pull request #35138 from jakevdp:pyrefly-fixes

Google-ML-Automation committed 101 days ago

985a077c

[checkify] delete non-hermetic test

mattjj committed 101 days ago

3d197e77

Use h100x2_full for the NVSHMEM test

apaszke committed 101 days ago

4aca5953

[Pallas] Use NumPy CPU random number generators instead of generating on TPU.

WindQAQ committed 101 days ago

d962a2bb

[Mosaic GPU] Add support for st.async (through FragmentedArray.store_tiled_async)

apaszke committed 102 days ago

115ceae8

[pallas:sc] The only supported array shapes now match physical vector register shapes

superbobry committed 102 days ago

26c69f1c

[Mosaic:GPU] Fix the cross-gpu barrier signal buffer sizes

PatriosTheGreat committed 102 days ago

6e6f672b

Newer Older