ROCm/jax

Pull Requests Commits

Update runner name in Bazel ROCm presubmit workflow

mkunredd committed 10 days ago

Verified 135e3f13

Add ROCm 10 wheel-version option for DPX smoke dispatches

Kunreddy committed 39 days ago

f2947fc7

pytest_rocm: add sandbox DPX runner for workflow_dispatch (#854)

mkunredd committed 44 days ago

Verified 955f99d3

Increase full ROCm pytest timeout to 240 minutes for multi-GPU DPX runs.

Kunreddy committed 45 days ago

fef60c88

Make run_pytest_rocm_smoke.sh executable for CI smoke runs.

Kunreddy committed 45 days ago

de6da288

Add ROCm pytest smoke mode for fast DPX validation (#864)

mkunredd committed 45 days ago

Verified f31bf287

ROCm pytest: 4 workers/GPU and DPX GPU preflight (#863)

mkunredd committed 46 days ago

Verified 814ea251

Route ROCm pytest to do-dpx-testing JAX runner labels (#862)

mkunredd committed 46 days ago

Verified c5330a61

Use sandbox DPX label for 8-GPU pytest runners (#856)

mkunredd committed 46 days ago

Verified 24ddec70

Route ROCm pytest to prod 1-GPU and DPX 4/8-GPU runners (#855)

mkunredd committed 46 days ago

Verified 6698570c

Bump pytest workers to 8/GPU, add bazel-skip toggle for nightly dispatch

magaonka-amd committed 52 days ago

7ed640bd

Route nightly ROCm pytest to the DPX node pool, halve workers/GPU

magaonka-amd committed 53 days ago

0a4040ea

[ROCm] Enable buffer_callback_test_gpu in ROCm CI

magaonka-amd committed 54 days ago

3aec0731

[ROCm] Enable the 35 device-agnostic bazel tests in ROCm CI

magaonka-amd committed 54 days ago

7d3bf91c

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

magaonka-amd committed 54 days ago

236b61fd

[ROCm] Tag ROCm CI job names with N / N+1 release tracks

magaonka-amd committed 54 days ago

9e0b44a8

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

mminutoli committed 54 days ago

b1daa4d7

Move Bazel test to the jax-base image (#831)

pelumi1163 committed 54 days ago

20a755ab

Enable additional single-GPU bazel tests in ROCm CI (#816)

magaonka-amd committed 54 days ago

98e67687

Merge pull request #40046 from jax-ml:postrelease

Google-ML-Automation committed 54 days ago

aa138ec5

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

Google-ML-Automation committed 54 days ago

83a74efd

[pallas:mgpu] Deprecated `plgpu.transpose_ref`

superbobry committed 54 days ago

468b30bb

[pallas] Moved `pl.reciprocal` into `pltpu`

superbobry committed 54 days ago

5d0f089f

Add deviceless PJRT topology serialization/deserialization and loaded executable version inspection in JAX.

brianseung-google committed 54 days ago

9fc1112c

Updates after release.

hawkinsp committed 54 days ago

5b9a6a4a

Merge v0.11.1 release branch.

hawkinsp committed 54 days ago

1776f2be

Add emit_pipeline primitives to the interpret mode.

IvyZX committed 54 days ago

c0215cbf

[NFC] Express `IsValidMmaTiling` using `OneOf` constraint

allanrenucci committed 54 days ago

0a96a9da

[Mosaic TPU] Support reshapes of form `(*src_major, src_minor) <-> (*dst_major, dst_minor)`, where `max(src_minor, dst_minor) <= 128` and one divides the other (row merging and splitting).

yueshengys committed 54 days ago

6f67db57

[pallas:mgpu] Removed `transform_ref`

superbobry committed 54 days ago

095198f8

Older