ROCm/jax

Pull Requests Commits

ScXfjiang committed 32 days ago

4c0559fd

[ROCm] Bump XLA module pin to gate collective kernels on has_peer_visible_atomics() (v0.11.1) (#876)

magaonka-amd committed 44 days ago

Verified fa79a53c

[ROCm] Bump XLA module pin to pick up the amd_smi migration (#872)

gulsumgudukbay committed 48 days ago

Verified 5e75202b

[ROCm] Bump XLA to pick up amd-smi migration (v0.11.1) (#867)

gulsumgudukbay committed 49 days ago

Verified 9606999e

[ROCm] Replace rocm-smi with amd-smi in CI diagnostics (v0.11.1) (#861)

gulsumgudukbay committed 50 days ago

Verified 6b928f60

Cherry-pick: [ROCm] Make test_vmap_ellipsis insensitive to reduced-precision matmuls (#853)

magaonka-amd committed 52 days ago

Verified 69d12681

Cherry-pick: [ROCm] Reduce pallas input/output aliasing test size on ROCm (#852)

magaonka-amd committed 57 days ago

Verified 2b611212

[ROCm] Give the pre-release branch a way of creating rc wheels.

mminutoli committed 59 days ago

937d6045

[ROCm] Manually bump of XLA pin.

mminutoli committed 59 days ago

a94e1bb9

[ROCm] nightly: build jax/jaxlib on fork runners

mminutoli committed 59 days ago

22d18faf

Prepare for JAX release 0.11.1

hawkinsp committed 59 days ago

2d666224

[XLA:GPU] Update lax test precision to account for numerics of new cuDNN frontend convolution pipeline.

derdrdirk committed 59 days ago

3ba3d761

Update XLA dependency to use revision

Google-ML-Automation committed 59 days ago

366ff357

Rename BUILD to BUILD.bazel in JAX toolchains.

Google-ML-Automation committed 59 days ago

ee8d0a1b

Merge pull request #40032 from hawkinsp:bazel

Google-ML-Automation committed 59 days ago

e309df58

Merge pull request #38946 from aybchan:aybchan/update-thor-gpuinfo-test

Google-ML-Automation committed 59 days ago

36cfce15

Update bazelversion to 7.7.1.

hawkinsp committed 59 days ago

3b1573a5

Merge pull request #40031 from jax-ml:fixeditable

Google-ML-Automation committed 59 days ago

91eac725

Removed `core_map_p` and everything which depended on it

superbobry committed 59 days ago

2dc5e1ab

Merge pull request #39759 from jakevdp:util-implicit-typevar

Google-ML-Automation committed 59 days ago

7be361ea

Merge pull request #40004 from jakevdp:api-typevar

Google-ML-Automation committed 59 days ago

9eabdd39

[Pallas TPU] Roll back exposing accumulator references to be used with explicit MXU APIs.

bchetioui committed 60 days ago

a66f606f

Update XLA dependency to use revision

Google-ML-Automation committed 60 days ago

2014e4c5

`core_map` now delegates to `mpmd_map`

superbobry committed 61 days ago

2fd39a4a

Support jax.debug.callback() in MPMD.

Google-ML-Automation committed 61 days ago

bcd95520

Merge pull request #40001 from ROCm:fix-rocm-wheel-summary-upstream

Google-ML-Automation committed 61 days ago

1e61f998

[Mosaic GPU] Reland back skipping multimem address initialization.

PatriosTheGreat committed 61 days ago

1b23d21c

Update XLA dependency to use revision

Google-ML-Automation committed 61 days ago

38f2e8d2

Fix zizmor suppressions for ml_actions images to use the correct format.

hawkinsp committed 61 days ago

aa28b5d6

Skip test that crashes under libtpu 0.0.46.

hawkinsp committed 61 days ago

a9fa9fbf

Older