ROCm/jax

Pull Requests Commits

Include prebuilt rocm plugin pkgs as deps to tests

alekstheod committed 11 days ago

600261dc

Properly mirror CompilationProviderOptions for jax

ermilovmaxim committed 23 days ago

e527f48c

Automated Code Change

Google-ML-Automation committed 23 days ago

1ab4f733

Rename to_cotangent_aval -> to_ct_aval. This is to match it to `to_ct_spec`.

yashk2810 committed 24 days ago

8d439f50

Merge pull request #35615 from mattjj:shmap-no-p

Google-ML-Automation committed 24 days ago

02519536

[hijax] generalize shmap closure handling for hitypes, add HipSpec.to_tangent_spec

mattjj committed 24 days ago

27094409

Reshape which unfolds the last dim into two will only be supported on TPU v5+ if the last dim is not divisible by 128.

yueshengys committed 24 days ago

bc94500c

Fix buggy unreduced add lowering.

mwhittaker committed 24 days ago

46f9bebc

Merge pull request #35581 from jakevdp:reduce-stack-frames

Google-ML-Automation committed 24 days ago

8b14bfd5

Add variable for rocm presubmit's to clone xla repo

ecalubaquib committed 24 days ago

c35c1ff6

[refactor] directly define array methods

jakevdp committed 24 days ago

93f24e8a

Reverts 1cd3ad579a92f2c25f3502b9d0683d5903afce78

WindQAQ committed 24 days ago

ca6fda28

Clean up some static initializer usage.

hawkinsp committed 24 days ago

42673a97

Jax updates for rules_python 1.8.4

ecalubaquib committed 24 days ago

fb41d650

Integrate LLVM at llvm/llvm-project@5ff5a1f14761

Google-ML-Automation committed 24 days ago

87db9ab6

[Mosaic TPU] Disallow DMAs with regular semaphores when src or tgt is in SMEM

apaszke committed 24 days ago

3564ddc9

[Mosaic] Add a `full_range` option to reciprocal.

WindQAQ committed 24 days ago

1cd3ad57

[Mosaic] Allow mask relayout with changed bitwidth and 1d tiling.

WindQAQ committed 24 days ago

9905f554

[pallas-triton] Explicitly disallow unsigned int operands in dot.

rdyro committed 24 days ago

9aa8b23b

[Mosaic:GPU] Rename is_comm_used as is_nvshmem_used

PatriosTheGreat committed 24 days ago

9fa9c3b7

[0.9.1] Update jaxlib version guards post-release.

danielsuo committed 24 days ago

47620776

[Mosaic GPU] Add support for f16/bf16 atomics

apaszke committed 24 days ago

7746f50f

[pallas:mgpu] Removed deprecated `Layout` aliases

superbobry committed 24 days ago

f37f484f

[Mosaic GPU] Refactor `FragmentedArray.broadcast` to use `broadcast_in_dim`.

allanrenucci committed 24 days ago

9e6b25ff

[Mosaic TPU] Stop using canonicalization rules for tiling propagation

apaszke committed 24 days ago

3562aa88

[NFC] Use `assert_never` for unreachable case in layout inference.

allanrenucci committed 24 days ago

40e91115

[Pallas][TPU kernel interpreter] Add optional logging to memory operations.

Google-ML-Automation committed 24 days ago

4ec06d48

[Mosaic GPU] Add basic support for atomic reductions while storing

apaszke committed 24 days ago

a9a9ce44

Merge pull request #35592 from superbobry:pyrefly

Google-ML-Automation committed 24 days ago

9ff85005

[Mosaic GPU][NFC] Add a `is_wg_semantics` method to `mosaic_gpu_test`.

dimitar-asenov committed 24 days ago

d36ae3ba

Older