ROCm/jax

Pull Requests Commits

Add pytest config

alekstheod committed 51 days ago

ef084e05

Switch to pytest instead of abseiltest in rocm

alekstheod committed 51 days ago

f3f7e763

alekstheod committed 52 days ago

94304906

Remove rbe sandboxing

alekstheod committed 52 days ago

10aad73f

Use init from rocm-jax

alekstheod committed 52 days ago

fe2cf6f6

alekstheod committed 52 days ago

22959814

Clean up filters multi-gpu single-gpu tests

alekstheod committed 53 days ago

57916539

Run only related tests

alekstheod committed 53 days ago

8e59fca7

Use option to mount volumes

alekstheod committed 53 days ago

5801656f

Trigger CI/CD pipelinee

alekstheod committed 53 days ago

5726ffe7

Use wheel or false to build jaxlib

alekstheod committed 53 days ago

a7be6ff6

Inject rocm plugin deps

alekstheod committed 53 days ago

fa4f298d

Migrate rocm-jax repo changes back to upstream jax

alekstheod committed 53 days ago

826d0e30

Add script for running bazel tests on ROCm

charleshofer committed 57 days ago

88276621

Move `experimental_compile_mpmd` from pathways to jaxlib.

emilyfertig committed 57 days ago

27274db2

Non-splat is unchanged

ImanHosseini committed 57 days ago

42c458cc

[Pallas] Shrink einshape generation size to avoid ASAN test timeouts

sharadmv committed 57 days ago

796d3bf8

Fix `_broadcast_in_dim_transpose_rule` to only consider unit_dims if the shape == 1 and spec is None. Fixes https://github.com/jax-ml/jax/issues/35181

yashk2810 committed 57 days ago

e69ef7b5

[Pallas/Mosaic GPU] Add cross-device barrier before running multi-device kernel with collective metadata

PatriosTheGreat committed 57 days ago

e5b542ac

[pallas] `mpmd_map` now supports lowering multiple kernels

superbobry committed 57 days ago

24c2f101

Replace the use of mlir::ExecutionEngine with a direct integration of llvm::orc::LLJIT for compiling and executing host-side LLVM IR.

khasanovaa committed 57 days ago

50daac0c

[CI] Update workflow files to use current best practices

MichaelHudgins committed 57 days ago

d7f29f75

[Mosaic] Allow int2 elementwise packing.

WindQAQ committed 58 days ago

afb5fba1

[Mosaic TPU] Add compatibility rules for the updated DenseElementsAttr encoding

apaszke committed 58 days ago

b78bd232

Merge pull request #34988 from gnecula:fix_pallas_vmap_shape_poly

Google-ML-Automation committed 58 days ago

62d36b59

[pallas:sc] Removed support for `pl.get_global`

superbobry committed 58 days ago

3c2d0db3

Fix for an incorrect LLVM integrate

apaszke committed 58 days ago

227c9849

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

Google-ML-Automation committed 58 days ago

aee33099

Add memory_space to typematch check

yashk2810 committed 58 days ago

540df654

Merge pull request #35095 from carlosgmartin:lax_reduce_doc_init_value_scalar

Google-ML-Automation committed 58 days ago

4a374b41

Older