ROCm/jax

Pull Requests Commits

Enable testSvdSubsetByIndex on ROCm

phambinhfin committed 1 day ago

0c98a018

[Mosaic GPU] Add support for all sensible conversions involving f8

apaszke committed 1 day ago

09e023e4

PallasLoweringDeterminismTest should only run on TPU backend

kanglant committed 1 day ago

9622105a

[Mosaic GPU] Use existing CUDA compute capability whenever available.

allanrenucci committed 1 day ago

30eaa17e

Add Bzlmod support for JAX.

meteorcloudy committed 1 day ago

228b61a7

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

Google-ML-Automation committed 1 day ago

ded22c03

Merge pull request #34586 from mattjj:config-tweak

Google-ML-Automation committed 1 day ago

359fa00e

tweak config for custom_vjp3: put in cache keys

mattjj committed 1 day ago

d87942f3

Check for args_not_mapped first and then fancy_primitive_batchers. If not, it leads to numerical issues (no idea why, but this was the behavior before skippable_batchers was deleted)

yashk2810 committed 1 day ago

768b1d6d

Merge pull request #34573 from jakevdp:refactor-hijax-tests

Google-ML-Automation committed 1 day ago

46d36125

Move the pallas_call batching logic which removes explicit_mesh_axis_names from AxisEnv into core.py and enter into that context in BatchTrace.process_primitive

yashk2810 committed 1 day ago

6fa4e041

[hijax,test] refactor hijax test suite

jakevdp committed 2 days ago

e131c5bb

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

Google-ML-Automation committed 2 days ago

b99ae997

ASAN support added to Hermetic C++

yuriivcs committed 2 days ago

d8cfcc39

Merge pull request #34449 from jakevdp:hijax-remat

Google-ML-Automation committed 2 days ago

3daa9845

[hijax] fix jit of remat of hijax primitive

jakevdp committed 2 days ago

da2785a3

Merge pull request #34541 from jakevdp:cond-remat

Google-ML-Automation committed 2 days ago

eda2370a

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

Google-ML-Automation committed 2 days ago

3cf4d328

[hijax] fix missing parameter in cond_p hijax lowering

jakevdp committed 2 days ago

b8f172f4

Add jax.tree.static()

jakevdp committed 2 days ago

4f3314e9

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

yashk2810 committed 2 days ago

db9bbb35

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

allanrenucci committed 2 days ago

47fd9f52

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

yashk2810 committed 2 days ago

fd12028f

Internal-only change.

jakeharmon8 committed 2 days ago

408d7cea

Use irfft to transpose rfft.

hawkinsp committed 2 days ago

182eab58

[Mosaic GPU][NFC] Style nits.

allanrenucci committed 2 days ago

419803e6

[Mosaic GPU] Generalize f8 casting routines

apaszke committed 2 days ago

0d2f49ab

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

allanrenucci committed 2 days ago

6eaada72

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

apaszke committed 2 days ago

17053a3f

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

bchetioui committed 2 days ago

1a0f79da

Older