ROCm/jax

Pull Requests Commits

Unskip testMultivariateNormalSingularCovariance on ROCm

AratiGanesh committed 22 days ago

fae7dd4b

Use irfft to transpose rfft.

hawkinsp committed 22 days ago

182eab58

[Mosaic GPU][NFC] Style nits.

allanrenucci committed 23 days ago

419803e6

[Mosaic GPU] Generalize f8 casting routines

apaszke committed 23 days ago

0d2f49ab

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

allanrenucci committed 23 days ago

6eaada72

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

apaszke committed 23 days ago

17053a3f

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

bchetioui committed 23 days ago

1a0f79da

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

Google-ML-Automation committed 23 days ago

c76c9e73

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

Google-ML-Automation committed 23 days ago

7b7e722a

[hijax] add a config for custom_vjp3

mattjj committed 23 days ago

2cdbff63

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

ybaturina committed 23 days ago

d83b0650

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

ZixuanJiang committed 23 days ago

e6b977c8

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

Google-ML-Automation committed 23 days ago

a7264587

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

emilyfertig committed 23 days ago

8fe97b1c

[Mosaic] Allow TC<->SC communication

vsytch committed 23 days ago

c8b43cd8

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

danielsuo committed 23 days ago

b8ecb137

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

Google-ML-Automation committed 23 days ago

67d98f30

[hijax] fix up last custom_vjp3 tests

mattjj committed 23 days ago

01c9adff

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

bchetioui committed 23 days ago

053846ca

Finish 0.9.0 release by removing dead code

yashk2810 committed 23 days ago

50734220

Decouple compilation and initialization in Mosaic GPU custom call.

allanrenucci committed 24 days ago

5199e5b0

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

golechwierowicz committed 24 days ago

b134a61a

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

apaszke committed 24 days ago

48205f93

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

apaszke committed 24 days ago

ecd37959

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

apaszke committed 24 days ago

b1760ad6

[Mosaic GPU] Take fewer arguments into MGPU kernels

apaszke committed 24 days ago

f2482d6d

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

bchetioui committed 24 days ago

9ff0b69a

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

allanrenucci committed 24 days ago

8f6a1989

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

bchetioui committed 24 days ago

feba25f7

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

Google-ML-Automation committed 24 days ago

c4481d42

Older