ROCm/jax

Pull Requests Commits

[ROCm] Enable ToeplitzSymmetricConstruction and condition number unit tests for ROCm.

tsrw2048 committed 6 hours ago

9c6368f4

[Mosaic GPU] Generalize f8 casting routines

apaszke committed 7 hours ago

0d2f49ab

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

allanrenucci committed 10 hours ago

6eaada72

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

apaszke committed 10 hours ago

17053a3f

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

bchetioui committed 12 hours ago

1a0f79da

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

Google-ML-Automation committed 13 hours ago

c76c9e73

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

Google-ML-Automation committed 16 hours ago

7b7e722a

[hijax] add a config for custom_vjp3

mattjj committed 17 hours ago

2cdbff63

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

ybaturina committed 20 hours ago

d83b0650

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

ZixuanJiang committed 21 hours ago

e6b977c8

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

Google-ML-Automation committed 22 hours ago

a7264587

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

emilyfertig committed 22 hours 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

[Mosaic GPU] Hand-roll a swizzling pattern to avoid bank conflicts in cross-warp reductions.

bchetioui committed 1 day ago

053846ca

Finish 0.9.0 release by removing dead code

yashk2810 committed 1 day ago

50734220

Decouple compilation and initialization in Mosaic GPU custom call.

allanrenucci committed 1 day ago

5199e5b0

[XLA:MGPU] Port ThreadIdxs to C++.

golechwierowicz committed 1 day ago

b134a61a

[Pallas:TPU] Add a two MXU test for the low-level MXU API

apaszke committed 1 day ago

48205f93

[Mosaic TPU] Allow more input types in the low-level MXU APIs

apaszke committed 1 day ago

ecd37959

[Pallas:TPU] Add a new low-level MXU API for TPUv7x

apaszke committed 1 day ago

b1760ad6

[Mosaic GPU] Take fewer arguments into MGPU kernels

apaszke committed 1 day ago

f2482d6d

[Pallas/Mosaic GPU] Disable Triton GEMMs in tests that do not use the autotuner.

bchetioui committed 1 day ago

9ff0b69a

[Mosaic GPU][NFC] Simplify expression with `utils.getelementptr`.

allanrenucci committed 1 day ago

8f6a1989

[Pallas/Mosaic GPU] Expand the docstring of `commit_smem` to encompass reads from SMEM.

bchetioui committed 1 day ago

feba25f7

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

Google-ML-Automation committed 1 day ago

c4481d42

Make auto_axes a no-op when the mesh is empty. This is required because library code can be written for no mesh context or a mesh context and auto-axes should work in both places. The other option is for user to do this kind of switching but I feel it's not too bad for JAX to do this since auto_axes already is a no-op if called under auto mode.

yashk2810 committed 1 day ago

ad30807b

Fix github lint breakage

yashk2810 committed 1 day ago

94ca6193

Older