ROCm/jax

Pull Requests Commits

[ROCm] Build jax/jaxlib against the TheRock 10.0 images

magaonka-amd committed 9 days ago

e156ed6b

[ROCm] Enable previously skipped pallas tests on gfx950

amd-jianli12 committed 9 days ago

0282fea3

[ROCm] Enable buffer_callback_test_gpu in ROCm CI

magaonka-amd committed 9 days ago

6b68e3a0

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

magaonka-amd committed 9 days ago

2d684a6e

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

magaonka-amd committed 9 days ago

9883fbb0

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

mminutoli committed 9 days ago

85e099b8

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

magaonka-amd committed 9 days ago

3d933e31

Merge pull request #40389 from ROCm:fix/rocm-multidim-real-fft-twiddle-cache

Google-ML-Automation committed 9 days ago

7d3d882c

Merge pull request #40370 from ROCm:user/mminutoli/drop_seccomp_unconfined_upstream

Google-ML-Automation committed 9 days ago

3662a88a

Make _arrays a read-only property of Array.

hawkinsp committed 10 days ago

39ec59b4

Don't barrier over literals everywhere in program_order

yashk2810 committed 10 days ago

cef889fc

Make Array.aval immutable.

hawkinsp committed 10 days ago

cd81d48f

[CI - Windows] Configure Bazel's JVM DNS to better survive short-lived lookup flakes.

belitskiy committed 10 days ago

620b2c55

Remove now-unused PyArray::ifrt_array() method.

hawkinsp committed 10 days ago

0636f0f2

[mgpu] Add `collective_axes` parameter to `try_cluster_cancel`.

chr1sj0nes committed 10 days ago

6f42033c

[Mosaic GPU] Fix missing `orders_tensor_cores=True` in test.

dimitar-asenov committed 10 days ago

f28bee43

Bump pinned XLA commit to get the LLVM integration patch fix.

belitskiy committed 10 days ago

bbb3fc3a

[Mosaic GPU] Replace NVVM ops with inline PTX in `dialect_lowering.py`.

bchetioui committed 10 days ago

1c7b170c

[Mosaic GPU] Add a `ScopedDiagnosticHandler` to propagate deserialization failures.

bchetioui committed 10 days ago

2334ef6d

[Mosaic GPU] Replace NVVM ops with inline PTX in `core.py`.

bchetioui committed 10 days ago

9ba3b04d

Merge pull request #40279 from saladpalad:mosaic-gpu-multidim-collective-commit

Google-ML-Automation committed 10 days ago

d0e27424

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

Google-ML-Automation committed 10 days ago

8126382e

[Mosaic GPU] Clean up `dialect_lowering`'s `LoweringContext` to always hold a `LaunchContext`.

bchetioui committed 10 days ago

2f838644

[ROCm] Work around a rocFFT twiddle cache bug in multi-dimensional real FFTs

magaonka-amd committed 10 days ago

9034a8a2

[Mosaic GPU] Emit a constant multicast mask for (2, 1, 1) collective commit arrivals

saladpalad committed 10 days ago

fc52bcfa

Replace TF_RETURN_IF_ERROR/TF_CHECK_OK with ABSL_RETURN_IF_ERROR/ABSL_CHECK_OK.

hawkinsp committed 10 days ago

b56c9a3a

[JAX] Remove _check_and_rearrange and skip_checks from PyArray constructor.

hawkinsp committed 10 days ago

f379ab77

[mgpu] Pipeline reductions across registers within thread.

chr1sj0nes committed 10 days ago

821f41c4

[Pallas] Update `CHANGELOG` following JAX's 0.11.1 release.

bchetioui committed 10 days ago

fc5da051

Replace uses of TF_ASSIGN_OR_RETURN with ABSL_ASSIGN_OR_RETURN.

hawkinsp committed 10 days ago

c5607f86

Older