ROCm/jax

Pull Requests Commits

Address review comments for approx_tanh ROCm implementation

phambinhfin committed 45 days ago

600fbd39

Implement approx_tanh for ROCm using OCML tanh function

phambinhfin committed 45 days ago

39ceb951

Add variable for rocm presubmit's to clone xla repo

ecalubaquib committed 45 days ago

c35c1ff6

Reverts 1cd3ad579a92f2c25f3502b9d0683d5903afce78

WindQAQ committed 45 days ago

ca6fda28

Clean up some static initializer usage.

hawkinsp committed 45 days ago

42673a97

Jax updates for rules_python 1.8.4

ecalubaquib committed 45 days ago

fb41d650

Integrate LLVM at llvm/llvm-project@5ff5a1f14761

Google-ML-Automation committed 45 days ago

87db9ab6

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

apaszke committed 45 days ago

3564ddc9

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

WindQAQ committed 45 days ago

1cd3ad57

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

WindQAQ committed 45 days ago

9905f554

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

rdyro committed 45 days ago

9aa8b23b

[Mosaic:GPU] Rename is_comm_used as is_nvshmem_used

PatriosTheGreat committed 45 days ago

9fa9c3b7

[0.9.1] Update jaxlib version guards post-release.

danielsuo committed 45 days ago

47620776

[Mosaic GPU] Add support for f16/bf16 atomics

apaszke committed 45 days ago

7746f50f

[pallas:mgpu] Removed deprecated `Layout` aliases

superbobry committed 45 days ago

f37f484f

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

allanrenucci committed 45 days ago

9e6b25ff

[Mosaic TPU] Stop using canonicalization rules for tiling propagation

apaszke committed 45 days ago

3562aa88

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

allanrenucci committed 45 days ago

40e91115

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

Google-ML-Automation committed 45 days ago

4ec06d48

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

apaszke committed 45 days ago

a9a9ce44

Merge pull request #35592 from superbobry:pyrefly

Google-ML-Automation committed 45 days ago

9ff85005

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

dimitar-asenov committed 45 days ago

d36ae3ba

Removed a few `type: ignore`s which were only necessary for mypy

superbobry committed 45 days ago

41f6727c

Merge pull request #35591 from superbobry:pyrefly

Google-ML-Automation committed 45 days ago

7ad75aca

Fix TensorStore implementation in Jax to use typed initialization.

Google-ML-Automation committed 45 days ago

1184dbaa

Removed unnecessary Pyrefly suppressions

superbobry committed 45 days ago

c307523f

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

Google-ML-Automation committed 45 days ago

e05dc3c4

Merge pull request #35587 from gnecula:poly_bug1

Google-ML-Automation committed 45 days ago

76cbef0e

[export] Add a test for shape polymorphism with invalid constraints

gnecula committed 45 days ago

a5ed9d25

Merge pull request #35576 from jakevdp:pyrefly-opt-einsum

Google-ML-Automation committed 45 days ago

42d97c87

Older