ROCm/jax

Pull Requests Commits

[ROCm] Enabled RNN unit test "test_no_workspace_overflow" for ROCm devices

tsrw2048 committed 4 hours ago

3b57b537

[Pallas/interpreter] Make the interpreter support barriers on GPU.

Google-ML-Automation committed 7 hours ago

b99ae997

ASAN support added to Hermetic C++

yuriivcs committed 7 hours ago

d8cfcc39

Merge pull request #34449 from jakevdp:hijax-remat

Google-ML-Automation committed 7 hours ago

3daa9845

[hijax] fix jit of remat of hijax primitive

jakevdp committed 8 hours ago

da2785a3

Merge pull request #34541 from jakevdp:cond-remat

Google-ML-Automation committed 8 hours ago

eda2370a

Merge pull request #34543 from jakevdp:jax-tree-static

Google-ML-Automation committed 8 hours ago

3cf4d328

[hijax] fix missing parameter in cond_p hijax lowering

jakevdp committed 8 hours ago

b8f172f4

Add jax.tree.static()

jakevdp committed 8 hours ago

4f3314e9

Add `_dot_general_batch_rule` back because some internal projects depend on it.

yashk2810 committed 8 hours ago

db9bbb35

[Mosaic GPU][NFC] Remove redundant calls to `opview`.

allanrenucci committed 8 hours ago

47fd9f52

Remove skippable_batchers and inline not mapped logic into primitive's batching rule

yashk2810 committed 9 hours ago

fd12028f

Internal-only change.

jakeharmon8 committed 9 hours ago

408d7cea

Use irfft to transpose rfft.

hawkinsp committed 9 hours ago

182eab58

[Mosaic GPU][NFC] Style nits.

allanrenucci committed 13 hours ago

419803e6

[Mosaic GPU] Generalize f8 casting routines

apaszke committed 15 hours ago

0d2f49ab

[Mosaic GPU] Run the kernel in `CustomCallTest.MosaicGpuUsesCommandBuffers`.

allanrenucci committed 18 hours ago

6eaada72

[Pallas:TPU] Add a matmul kernel test using the low-level MXU APIs

apaszke committed 18 hours ago

17053a3f

[Mosaic GPU] Leverage additional scratch space in cross-warp reductions.

bchetioui committed 20 hours ago

1a0f79da

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

Google-ML-Automation committed 21 hours ago

c76c9e73

Merge pull request #34548 from mattjj:custom-vjp3-6

Google-ML-Automation committed 23 hours ago

7b7e722a

[hijax] add a config for custom_vjp3

mattjj committed 1 day ago

2cdbff63

Update the requirements' lock files post the JAX 0.9.0 release.

ybaturina committed 1 day ago

d83b0650

Disable the `contracting2` test case from `test_dot_general`.

ZixuanJiang committed 1 day ago

e6b977c8

Merge pull request #34264 from emilyfertig:emilyaf-xhost-device-put-docs

Google-ML-Automation committed 1 day ago

a7264587

Document cross-host `device_put` and running on a subset of devices.

emilyfertig committed 1 day ago

8fe97b1c

[Mosaic] Allow TC<->SC communication

vsytch committed 1 day ago

c8b43cd8

[pmap] Inline `multihost_utils` `impl` functions for host_local <> global conversions, `_shared_code_pmap`, and `_prepare_pmap`.

danielsuo committed 1 day ago

b8ecb137

Merge pull request #34423 from mattjj:custom-vjp3-youandme-5

Google-ML-Automation committed 1 day ago

67d98f30

[hijax] fix up last custom_vjp3 tests

mattjj committed 1 day ago

01c9adff

Older