Commits ROCm/jax - SemanticDiff

[ROCm] Migrate CI runners to Digital Ocean MI350X

mminutoli committed 137 days ago

f257cd51

[ROCm] Put TheRock SDK bin on PATH via so rocminfo/rocm-smi resolve in /opt/rocm-less images

gulsumgudukbay committed 137 days ago

b8c52b15

[ROCm] Set HSA_NO_SCRATCH_RECLAIM=1 in ROCm test workflows

gulsumgudukbay committed 137 days ago

4a864c75

[ROCm] Embed TheRock-aware wheel RUNPATH at link time

gulsumgudukbay committed 137 days ago

e75909f5

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

mminutoli committed 137 days ago

114daa22

[Mosaic GPU] Add transform commutation for cluster_ref.

jntang-google committed 137 days ago

a85b98a6

Allow control over jit's cpp cache.

yashk2810 committed 137 days ago

d6485690

[Pallas/Mosaic GPU][Docs] Add hyperlinks to various API bits throughout the reference.

bchetioui committed 137 days ago

2b684b82

[Pallas:MGPU] Support batch dimension for TMEM references

allanrenucci committed 137 days ago

c696f3a8

[Mosaic:GPU] Move gpu_distributed test from NVSHMEM to collective metadata.

PatriosTheGreat committed 137 days ago

c98d8fa0

Fix numerical flakiness in StateHypothesisTest.test_addupdate_vmap on GPU.

hawkinsp committed 137 days ago

dbf9fb12

Disable some flaky tests under sanitizer builds.

hawkinsp committed 137 days ago

9663628c

Allow test_dot_ag_pipeline_grad to run on all devices

yashk2810 committed 137 days ago

ecf86679

Update xla_transform_test to unregister callbacks after each unit test.

Google-ML-Automation committed 137 days ago

e7316c20

Only use `compute_on2` in `test_dot_ag_pipeline_grad_tpu` on TPU 7x

yashk2810 committed 137 days ago

e75ec61e

Collapse Bazel output by default in Github Actions.

hawkinsp committed 137 days ago

7ed89289

[pallas:mgpu] Check that at most one of `out_shape`/`out_type` and `scratch_shapes`/`scratch_types` is provided

superbobry committed 137 days ago

c1fa7237

Reverts aa296d2a281ded40c47b68cf695432d6e0c793ad

superbobry committed 138 days ago

30db8844

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

Google-ML-Automation committed 138 days ago

db6acb79

[consts] Fix tests for JAX_USE_SIMPLIFIED_JAXPR_CONSTANTS=1

gnecula committed 138 days ago

8f2f1a8b

remove a print and delete the GPU test (merged with TPU test now)

yashk2810 committed 138 days ago

7267cb12

Add an example for pipelining FSDP with scan + custom_vjp (just with 1 dot + 1 ag on fwd).

yashk2810 committed 138 days ago

9e3199ba

Update xla_transform_test to unregister callbacks after each unit test.

Google-ML-Automation committed 138 days ago

a0bbbc1e

jax.experimental.mosaic: expose a few more tpu_custom_call APIs

Jake VanderPlas committed 138 days ago

39b296a0

Merge pull request #37888 from mattjj:remat3-scan

Google-ML-Automation committed 138 days ago

bccacd50

[remat3] add scan/cond/custom_jvp/vjp rules, tests, many small fixes

mattjj committed 138 days ago

bf1e6d9b

Merge pull request #38372 from ROCm:fix/pallas-fused-attention-bwd-resource-exhausted-skip-upstream

Google-ML-Automation committed 138 days ago

0dfebd20

Rework the expectations for tpu_pallas_memory_space_test so that assertions, about memory space coloring, are made on the initial HLO passed to the compiler, not on the optimized HLO, just before lowering.

sparc1998 committed 138 days ago

71fcd980

[NFC] Fix weird formatting in `mosaic_gpu_test.py`.

allanrenucci committed 138 days ago

68be44eb

Add wrap_negative_indices to take_along_axis and avoid upcasting the take_along_dim index when wrap_negative_indices=False.

18praveenb committed 138 days ago

3c174b94