ROCm/jax

Pull Requests Commits

pytest_rocm: add 4-GPU and 8-GPU DPX runner options

mkunredd committed 5 days ago

4b1c2e86

pytest_rocm: increase timeout to 240min for DPX (slower partition)

mkunredd committed 5 days ago

17b95503

pytest_rocm: add dpx dind runner option

mkunredd committed 6 days ago

d216cea0

Add 1-GPU isolated nodind label for manual pytest validation

Kunreddy committed 11 days ago

6fba2db2

Expose nodind runner label for manual ROCm pytest validation

Kunreddy committed 11 days ago

b32e5cbc

Mount /home/runner/_work in ROCm workflow containers

Kunreddy committed 11 days ago

0c275445

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

mminutoli committed 11 days ago

d8b3f548

Remove stale ROCm skips from four GPU test cases

magaonka-amd committed 11 days ago

839c6488

Enable 12 more single-GPU bazel test targets in ROCm CI

magaonka-amd committed 11 days ago

e9978550

[ROCm] Enable previously skipped unit tests on ROCm platform

magaonka-amd committed 11 days ago

bde178f3

Simplify pytest parallelism calculation for ROCm CI

magaonka-amd committed 11 days ago

a0587192

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

pelumi1163 committed 11 days ago

b22e303b

ci: retire ROCm 7.2.0 and TheRock 7.13 flows; run only TheRock 7.14 + latest

magaonka-amd committed 11 days ago

398f6bfb

ci: bump ROCm release wheel tests to TheRock 7.14 GA (#817)

Ruturaj4 committed 11 days ago

39bcc56e

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

magaonka-amd committed 11 days ago

65aa7f4b

ci: run ROCm release wheel tests on TheRock 7.14.0rc3 images (#812)

Ruturaj4 committed 11 days ago

1247cc28

[ROCm] Add release pipeline CI workflows

mminutoli committed 11 days ago

18402cf5

[pallas:mgpu] Automatically infer `PROMISE_IN_BOUNDS` in `emit_pipeline`

superbobry committed 11 days ago

4861a40d

Replace closed form expansion of 2x2 and 3x3 determinants with a closed form LU decomposition with row pivoting.

hawkinsp committed 11 days ago

a25d587e

[triton] Added support for host-side TMA descriptors

superbobry committed 11 days ago

8585ff73

[mosaic_gpu] The `barrier` operand to `async_load` is now optional on pre-Hopper GPUs

superbobry committed 11 days ago

263eed93

Implement AOT compilation for autotuned JAX-Triton kernels.

EusebioDM committed 11 days ago

8e07bd40

Merge pull request #39895 from jax-ml:bzlmod

Google-ML-Automation committed 11 days ago

51b5b00f

[pallas] Deprecated `pl.core_map` in favor of `pl.kernel`

superbobry committed 12 days ago

d3287b70

Switch to bzlmod by default.

hawkinsp committed 12 days ago

aeff4e4b

Merge pull request #39901 from jax-ml:reqs

Google-ML-Automation committed 12 days ago

24bd68dd

Regenerate requirements lock files.

hawkinsp committed 12 days ago

aa742407

Add emulation for int4 cmpi

Sam Ginzburg committed 12 days ago

798b72fa

Merge pull request #39915 from mattjj:push-ustwymvzzpon

Google-ML-Automation committed 12 days ago

da2d20de

Make dce_sink a VJPHiPrimitive so that it works with sharding

yashk2810 committed 12 days ago

4b315c97

Older