ROCm/jax

Pull Requests Commits

Restore local sysroot option

alekstheod committed 5 hours ago

6ffc177c

Switch to build jax using hermetic sysroot

alekstheod committed 5 hours ago

8d8b6b75

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

Google-ML-Automation committed 1 day 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 2 days 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 3 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 3 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 3 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 4 days ago

17a4531e

Merge pull request #37927 from gnecula:clean_jax2tf

Google-ML-Automation committed 4 days ago

5e6b7ccc

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

yashk2810 committed 4 days ago

2c128ca1

Add ur rule for select_n_p

yashk2810 committed 4 days ago

8eaeabb0

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

Google-ML-Automation committed 4 days ago

0abeb0ab

Add tests for uneven axis index groups in collectives.

Google-ML-Automation committed 4 days ago

2d409b85

Print jaxprs with >= 10 equations out of line.

hawkinsp committed 4 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

[Mosaic] Support unpack_elementwise to unpack to non-32-bit dtypes.

WindQAQ committed 4 days ago

a297191a

Document strictly-increasing requirement for RegularGridInterpolator without runtime validation

adityasingh2400 committed 4 days ago

8d62a2f2

Reverts 22781e60b16969a6594e86a2f678848b365ede1d

hawkinsp committed 4 days ago

ac12cbaa

Merge pull request #37928 from lanarkite99:resume/jax-bcoo-vmap-grad

Google-ML-Automation committed 4 days ago

0a98102a

Merge pull request #37806 from samanklesaria:add_f_cast_check

Google-ML-Automation committed 4 days ago

263d192b

Older