ROCm/jax

Pull Requests Commits

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

mkunredd committed 55 days ago

4b1c2e86

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

mkunredd committed 55 days ago

17b95503

pytest_rocm: add dpx dind runner option

mkunredd committed 56 days ago

d216cea0

Add 1-GPU isolated nodind label for manual pytest validation

Kunreddy committed 61 days ago

6fba2db2

Expose nodind runner label for manual ROCm pytest validation

Kunreddy committed 61 days ago

b32e5cbc

Mount /home/runner/_work in ROCm workflow containers

Kunreddy committed 61 days ago

0c275445

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

mminutoli committed 61 days ago

d8b3f548

Remove stale ROCm skips from four GPU test cases

magaonka-amd committed 61 days ago

839c6488

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

magaonka-amd committed 61 days ago

e9978550

[ROCm] Enable previously skipped unit tests on ROCm platform

magaonka-amd committed 61 days ago

bde178f3

Simplify pytest parallelism calculation for ROCm CI

magaonka-amd committed 61 days ago

a0587192

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

pelumi1163 committed 61 days ago

b22e303b

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

magaonka-amd committed 61 days ago

398f6bfb

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

Ruturaj4 committed 61 days ago

39bcc56e

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

magaonka-amd committed 61 days ago

65aa7f4b

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

Ruturaj4 committed 61 days ago

1247cc28

[ROCm] Add release pipeline CI workflows

mminutoli committed 61 days ago

18402cf5

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

superbobry committed 61 days ago

4861a40d

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

hawkinsp committed 61 days ago

a25d587e

[triton] Added support for host-side TMA descriptors

superbobry committed 61 days ago

8585ff73

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

superbobry committed 62 days ago

263eed93

Implement AOT compilation for autotuned JAX-Triton kernels.

EusebioDM committed 62 days ago

8e07bd40

Merge pull request #39895 from jax-ml:bzlmod

Google-ML-Automation committed 62 days ago

51b5b00f

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

superbobry committed 62 days ago

d3287b70

Switch to bzlmod by default.

hawkinsp committed 62 days ago

aeff4e4b

Merge pull request #39901 from jax-ml:reqs

Google-ML-Automation committed 62 days ago

24bd68dd

Regenerate requirements lock files.

hawkinsp committed 62 days ago

aa742407

Add emulation for int4 cmpi

SamGinzburg committed 62 days ago

798b72fa

Merge pull request #39915 from mattjj:push-ustwymvzzpon

Google-ML-Automation committed 62 days ago

da2d20de

Make dce_sink a VJPHiPrimitive so that it works with sharding

yashk2810 committed 62 days ago

4b315c97

Older