Commits ROCm/jax - SemanticDiff

[ROCm] Added a release validation workflow for JAX ROCm wheels.

tsrw2048 committed 33 days ago

3a524d46

Add org-based guards to all CI caller workflows

mminutoli committed 33 days ago

472e57ed

[ROCm] Version-tier S3 layout, standardized CI

mminutoli committed 33 days ago

bbfc0f9d

[pallas/mgpu] Add a copy_smem_to_gmem scatter operation.

cperivol committed 33 days ago

45e42c27

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

cperivol committed 33 days ago

3926ae67

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

Google-ML-Automation committed 33 days ago

f056bd87

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

gnecula committed 34 days ago

969f945d

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

Google-ML-Automation committed 34 days ago

1542dd1c

Reverts ac12cbaa1f877ef38a128000b8c561f136fc7f01

hawkinsp committed 35 days ago

82c22a11

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

Google-ML-Automation committed 35 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 36 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 36 days ago

0e68e5f4

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

Jake VanderPlas committed 36 days ago

84624e9a

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

yashk2810 committed 36 days ago

1c32cfdc

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

Google-ML-Automation committed 36 days ago

17a4531e

Merge pull request #37927 from gnecula:clean_jax2tf

Google-ML-Automation committed 37 days ago

5e6b7ccc

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

yashk2810 committed 37 days ago

2c128ca1

Add ur rule for select_n_p

yashk2810 committed 37 days ago

8eaeabb0

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

Google-ML-Automation committed 37 days ago

0abeb0ab

Add tests for uneven axis index groups in collectives.

Google-ML-Automation committed 37 days ago

2d409b85

Print jaxprs with >= 10 equations out of line.

hawkinsp committed 37 days ago

efcd4395

Add tests for uneven axis index groups in collectives.

hawkinsp committed 37 days ago

8ff771c7

Merge pull request #37798 from samanklesaria:t_dist_type_fix

Google-ML-Automation committed 37 days ago

e96d8284

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

Google-ML-Automation committed 37 days ago

18e3b114

Merge pull request #37759 from samanklesaria:beta_arg_check

Google-ML-Automation committed 37 days ago

00380b7e

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

Google-ML-Automation committed 37 days ago

2750d75a

Remove stale version tests.

hawkinsp committed 37 days ago

08aa8996

Merge pull request #37800 from samanklesaria:pareto_dtype_fix

Google-ML-Automation committed 37 days ago

f4829185

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

WindQAQ committed 37 days ago

a297191a

Document strictly-increasing requirement for RegularGridInterpolator without runtime validation

adityasingh2400 committed 37 days ago

8d62a2f2