ROCm/jax

Pull Requests Commits

Route ROCm pytest to do-dpx-testing JAX runner labels

Kunreddy committed 2 days ago

05e35211

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

mkunredd committed 2 days ago

Verified 24ddec70

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

mkunredd committed 2 days ago

Verified 6698570c

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

magaonka-amd committed 8 days ago

7ed640bd

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

magaonka-amd committed 9 days ago

0a4040ea

[ROCm] Enable buffer_callback_test_gpu in ROCm CI

magaonka-amd committed 10 days ago

3aec0731

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

magaonka-amd committed 10 days ago

7d3bf91c

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

magaonka-amd committed 10 days ago

236b61fd

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

magaonka-amd committed 10 days ago

9e0b44a8

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

mminutoli committed 10 days ago

b1daa4d7

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

pelumi1163 committed 10 days ago

20a755ab

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

magaonka-amd committed 10 days ago

98e67687

Merge pull request #40046 from jax-ml:postrelease

Google-ML-Automation committed 10 days ago

aa138ec5

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

Google-ML-Automation committed 10 days ago

83a74efd

[pallas:mgpu] Deprecated `plgpu.transpose_ref`

superbobry committed 10 days ago

468b30bb

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

superbobry committed 10 days ago

5d0f089f

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

brianseung-google committed 10 days ago

9fc1112c

Updates after release.

hawkinsp committed 10 days ago

5b9a6a4a

Merge v0.11.1 release branch.

hawkinsp committed 10 days ago

1776f2be

Add emit_pipeline primitives to the interpret mode.

IvyZX committed 10 days ago

c0215cbf

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

allanrenucci committed 10 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 10 days ago

6f67db57

[pallas:mgpu] Removed `transform_ref`

superbobry committed 10 days ago

095198f8

[pallas] Removed `MemoryRefBase`

superbobry committed 10 days ago

f1d7d865

PR #40041: Avoid quadratic tracing behaviour

olupton committed 10 days ago

45962b55

[typing] Use implicit TypeVar in jax.lax sources

Jake VanderPlas committed 10 days ago

9e57d0b1

Prepare for JAX release 0.11.1

hawkinsp committed 10 days ago

2d666224

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

derdrdirk committed 10 days ago

3ba3d761

Update XLA dependency to use revision

Google-ML-Automation committed 10 days ago

366ff357

Rename BUILD to BUILD.bazel in JAX toolchains.

Google-ML-Automation committed 10 days ago

ee8d0a1b

Older