ROCm/jax

Pull Requests Commits

Skip test_tridiagonal_solve on ROCm

AratiGanesh committed 21 days ago

50ad5ba9

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

yashk2810 committed 21 days ago

fd12028f

Internal-only change.

jakeharmon8 committed 21 days ago

408d7cea

Use irfft to transpose rfft.

hawkinsp committed 21 days ago

182eab58

[Mosaic GPU][NFC] Style nits.

allanrenucci committed 21 days ago

419803e6

[Mosaic GPU] Generalize f8 casting routines

apaszke committed 21 days ago

0d2f49ab

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

allanrenucci committed 21 days ago

6eaada72

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

apaszke committed 21 days ago

17053a3f

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

bchetioui committed 21 days ago

1a0f79da

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

Google-ML-Automation committed 22 days ago

c76c9e73

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

Google-ML-Automation committed 22 days ago

7b7e722a

[hijax] add a config for custom_vjp3

mattjj committed 22 days ago

2cdbff63

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

ybaturina committed 22 days ago

d83b0650

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

ZixuanJiang committed 22 days ago

e6b977c8

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

Google-ML-Automation committed 22 days ago

a7264587

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

emilyfertig committed 22 days ago

8fe97b1c

[Mosaic] Allow TC<->SC communication

vsytch committed 22 days ago

c8b43cd8

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

danielsuo committed 22 days ago

b8ecb137

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

Google-ML-Automation committed 22 days ago

67d98f30

[hijax] fix up last custom_vjp3 tests

mattjj committed 22 days ago

01c9adff

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

bchetioui committed 22 days ago

053846ca

Finish 0.9.0 release by removing dead code

yashk2810 committed 22 days ago

50734220

Decouple compilation and initialization in Mosaic GPU custom call.

allanrenucci committed 22 days ago

5199e5b0

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

golechwierowicz committed 22 days ago

b134a61a

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

apaszke committed 22 days ago

48205f93

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

apaszke committed 22 days ago

ecd37959

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

apaszke committed 22 days ago

b1760ad6

[Mosaic GPU] Take fewer arguments into MGPU kernels

apaszke committed 22 days ago

f2482d6d

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

bchetioui committed 22 days ago

9ff0b69a

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

allanrenucci committed 22 days ago

8f6a1989

Older