ROCm/jax

Pull Requests Commits

[ROCm] Add TheRock 7.13 and TheRock latest to CI test matrices

mminutoli committed 1 hour ago

b5d7662c

[ROCm] Version-tier S3 layout, standardized CI

mminutoli committed 1 hour ago

ba1ddd39

Add org-based guards to all CI caller workflows

mminutoli committed 1 hour ago

7df00da3

Reverts 82c22a11000c036dbcf5a662b94a74ccc11d02bb

hawkinsp committed 2 hours ago

6305fc44

Add support for passing the grid indices to cached Mosaic lowerings.

hawkinsp committed 2 hours ago

228b9345

[pallas/mgpu] Add a copy_smem_to_gmem scatter operation.

cperivol committed 6 hours ago

45e42c27

[pallas:mgpu] Implement shift left and shift right operators.

cperivol committed 6 hours ago

3926ae67

Update XLA dependency to use revision http://github.com/openxla/xla/commit/0ef91e2443e93d9d4649c1a73bdd71f95319b12c

Google-ML-Automation committed 12 hours ago

f056bd87

[pallas] Fix Pallas Mosaic 64-bit mode loop indexing

gnecula committed 1 day ago

969f945d

Update XLA dependency to use revision http://github.com/openxla/xla/commit/93f41f404fa2e38df46cd8cf8220613ca5182c48

Google-ML-Automation committed 1 day ago

1542dd1c

Reverts ac12cbaa1f877ef38a128000b8c561f136fc7f01

hawkinsp committed 2 days ago

82c22a11

Update XLA dependency to use revision http://github.com/openxla/xla/commit/9398ef15bc0d6c62d72236db7ea9c7b451dfe51e

Google-ML-Automation committed 2 days ago

12350eea

Add a test to make sure that if input is unreduced on fwd (with check_vma=False on shard_map), shard_map inserts a psum on bwd similar to invarying.

yashk2810 committed 2 days ago

636a1a98

Use the cotangent spec when calculating which mesh axes to do psum over if check_vma=False in shard_map transpose.

yashk2810 committed 2 days ago

0e68e5f4

JAX: Add `mulhi` primitive which lowers to the recently-added `chlo.mulhi` op.

Jake VanderPlas committed 3 days ago

84624e9a

Add a test for remat transpose when the primal inputs are reduced. Make sure the cotangents are unreduced

yashk2810 committed 3 days ago

1c32cfdc

Update XLA dependency to use revision http://github.com/openxla/xla/commit/4ef447c041b91498bf04b30e9395f703facff9cf

Google-ML-Automation committed 3 days ago

17a4531e

Merge pull request #37927 from gnecula:clean_jax2tf

Google-ML-Automation committed 3 days ago

5e6b7ccc

Make sure out_cts passed to remat ad.backward_pass are to_ct_aval()

yashk2810 committed 3 days ago

2c128ca1

Add ur rule for select_n_p

yashk2810 committed 3 days ago

8eaeabb0

PR #35699: Strip Mosaic kernel debug info and exclude autotune cache dir from cache keys

Google-ML-Automation committed 3 days ago

0abeb0ab

Add tests for uneven axis index groups in collectives.

Google-ML-Automation committed 3 days ago

2d409b85

Print jaxprs with >= 10 equations out of line.

hawkinsp committed 3 days ago

efcd4395

Add tests for uneven axis index groups in collectives.

hawkinsp committed 4 days ago

8ff771c7

Merge pull request #37798 from samanklesaria:t_dist_type_fix

Google-ML-Automation committed 4 days ago

e96d8284

Merge pull request #37916 from adityasingh2400:doc-regulargridinterpolator-strictly-increasing-36499

Google-ML-Automation committed 4 days ago

18e3b114

Merge pull request #37759 from samanklesaria:beta_arg_check

Google-ML-Automation committed 4 days ago

00380b7e

Merge pull request #35699 from colin-ant:compilation-cache-6mar

Google-ML-Automation committed 4 days ago

2750d75a

Remove stale version tests.

hawkinsp committed 4 days ago

08aa8996

Merge pull request #37800 from samanklesaria:pareto_dtype_fix

Google-ML-Automation committed 4 days ago

f4829185

Older