Rocm jaxlib v0.9.0#671

Open

Ruturaj4 wants to merge 6584 commits intomainfrom

rocm-jaxlib-v0.9.0

This pull request is big! We're only showing the most recent 250 commits

Commits on Jan 7, 2026

Merge pull request #34204 from jakevdp:doc-sidebar
Google-ML-Automation
committed
Merge pull request #32268 from samanklesaria:issues/32267
Google-ML-Automation
committed
Prefer isinstance(x, type) over type.isinstance

boomanaiden154
authored and
Google-ML-Automation
committed
Add jax.experimental.random to the wheel build
jakevdp
committed
fix shard_map transpose explicit sharding zero unsharding bug

mattjj
and
yashk2810
committed
add optional `explain` callback for weakref_lru_cache misses

authored and
Google-ML-Automation
committed
Merge pull request #34209 from jakevdp:fix-wheel
Google-ML-Automation
committed

Commits on Jan 8, 2026

Commits on Jan 9, 2026

Commits on Jan 10, 2026

Commits on Jan 11, 2026

Commits on Jan 12, 2026

Commits on Jan 13, 2026

Commits on Jan 14, 2026

Commits on Jan 15, 2026

Commits on Jan 16, 2026

Bump XLA version.

mwhittaker
authored and
Google-ML-Automation
committed
Prepare for JAX release 0.9.0
mwhittaker
committed
Increase shard count for TPU and GPU tests back to 5 for api_test.py.

belitskiy
authored and
mwhittaker
committed
Fix wheel sources tests for Windows platform.

ybaturina
authored and
mwhittaker
committed
Disable `tests/multiprocess:socket_transfer_test` since it's failing flakily in all backends.

emilyfertig
authored and
mwhittaker
committed
Re-enable `socket_transfer_test` internally.

emilyfertig
authored and
mwhittaker
committed
Add libtpu date guard for failing test.
mwhittaker
committed
Skip tpu_pallas_distributed_test on 7x.
mwhittaker
committed
Add libtpu guard to failing tpu_trace_value_test.
mwhittaker
committed
Disable `tpu_splash_attention_kernel_test` on TPU v7x.
mwhittaker
committed
Remove failing test_itof_dot_canonicalization_fails_without_compat_mode test.

18praveenb
authored and
mwhittaker
committed

Commits on Jan 20, 2026

Use maxsize=None with trace_to_jaxpr's weakref_lru_cache to get more cache hits. This also restores previous behavior of lu.cache.

yashk2810
authored and
mwhittaker
committed

Commits on Jan 23, 2026

Remove nvidia_wheel_versions

charleshofer
authored and
Ruturaj4
committed
Make jaxlib targets visible

charleshofer
authored and
Ruturaj4
committed
hipblas typedef fix

charleshofer
authored and
Ruturaj4
committed
No GPU fail

charleshofer
authored and
Ruturaj4
committed
Wrap HIP inline functions in anonymous namespaces in vendor.h

AratiGanesh
authored and
Ruturaj4
committed
SWDEV-512768 - Replace hipGetLastError with hipExtGetLastError

dsicarov-amd
authored and
Ruturaj4
committed
Add shared utility function get_rocm_version to test_util.py

charleshofer
authored and
Ruturaj4
committed
Fix hipSparse CSR algorithm mappings for ROCm 7

phambinhfin
authored and
Ruturaj4
committed
Make nvidia version data optional for ROCm builds

phambinhfin
authored and
Ruturaj4
committed
Fix v_pages quantization and adjust test params for ROCm compatibilit… (#560)

phambinhfin
authored and
Ruturaj4
committed
Address LLVM assertion failure due to a multithreaded use. Update .gitignore (#563)

Arech8
authored and
Ruturaj4
committed
Add skip of test_is_finite() on Cuda (#565)

Arech8
authored and
Ruturaj4
committed
Add rocm test requirements file (#570)

AratiGanesh
authored and
Ruturaj4
committed
Let the unit tests use build.py for setting up Bazel commands for unit tests (#582)

charleshofer
authored and
Ruturaj4
committed
adding abort logic to rocm/jax (#590)

gulsumgudukbay
authored and
Ruturaj4
committed
Skip is_finite tests on ROCm (not in Triton lowering for jax 0.8.0) (#597)

phambinhfin
authored and
Ruturaj4
committed
Fix shared memory limit check for ROCm in test_dot (#596)

phambinhfin
authored and
Ruturaj4
committed
Fix Numpy signatures test (#598)

authored and
Ruturaj4
committed
Fix GPU lowering rule for SVD on ROCm devices (#600)

tsrw2048
authored and
Ruturaj4
committed
Fixed merge conflicts when moving GESVDJ commit from JAX 0.8.0 to 0.8.2 (#602)

tsrw2048
authored and
Ruturaj4
committed
fix merge arts
Ruturaj4
committed
Enabled testTridiagonal to run on ROCm devices (#607)

tsrw2048
authored and
Ruturaj4
committed
Enabled ToeplitzSymmetricConstruction unit tests for ROCm devices (#608)

tsrw2048
authored and
Ruturaj4
committed
Enable testSvdSubsetByIndex for ROCm with subset_by_index skip (#603)

phambinhfin
authored and
Ruturaj4
committed
Fix KeyError for bytes_reservable_limit on ROCm

magaonka-amd
authored and
Ruturaj4
committed
Backport cuda_array_interface , testDotAlgorithm, Optimizer test fixes from rocm-jaxlib-v0.8.0 (#635)

magaonka-amd
authored and
Ruturaj4
committed
Enable RngShardingTests (#644)

gulsumgudukbay
authored and
Ruturaj4
committed
Enable reduce_window tests on ROCm (#643)

magaonka-amd
authored and
Ruturaj4
committed
Enable test_variadic_reduce_window on ROCm (#647)

magaonka-amd
authored and
Ruturaj4
committed
Unskip supported dtypes for testConvolutionsPreferredElementType (#649)

gulsumgudukbay
authored and
Ruturaj4
committed
Enabled test for condition number on ROCm devices. (#613)

tsrw2048
authored and
Ruturaj4
committed
Enabled RNN unit test: test_no_workspace_overflow for ROCm devices (#624)

tsrw2048
authored and
Ruturaj4
committed
Added changes from PR #626 and PR #645. This also fixes merge conflicts from upstream.

tsrw2048
authored and
Ruturaj4
committed

Commits on Jan 29, 2026

[Pallas] Fix ROCm GPU architecture detection and route to Triton backend
Ruturaj4
committed
Enable array interoperability tests on ROCm platform (#660)

magaonka-amd
authored and
Ruturaj4
committed
Skip sparse tests on ROCm due to hipSPARSE issue (#652)

magaonka-amd
authored and
Ruturaj4
committed
Update sparse test skip messages in v0.8.2 (#653)

magaonka-amd
authored and
Ruturaj4
committed
Port skip for "test_prim_tridiagonal_solve" tests from JAX v0.8.0 to v0.8.2 (#658)

tsrw2048
authored and
Ruturaj4
committed
Fix test_cuda_array_interface test skip condition (#657)

magaonka-amd
authored and
Ruturaj4
committed
Enable testMultivariateNormalSingularCovariance on ROCm (#666)

AratiGanesh
authored and
Ruturaj4
committed
Skip test_batch_axis_sharding_jvp because of hipSPARSE issue (#667)

AratiGanesh
authored and
Ruturaj4
committed
Skip test_tridiagonal_solve on ROCm due to hipSPARSE numerical errors (#668)

AratiGanesh
authored and
Ruturaj4
committed
Update Skip Reason Outputs (#663)

gulsumgudukbay
authored and
Ruturaj4
committed
Skip testCudaArrayInterfaceOnNonCudaFails on ROCm platform (#677)
magaonka-amd
authored

Commits on Feb 3, 2026

Enable lobpcg tests on ROCm platform (#681)
magaonka-amd
authored

Commits on Feb 4, 2026

Enable lax backend scipy tests on ROCm GPUs (#687)
magaonka-amd
authored

Commits on Feb 5, 2026

Add ROCm encoding for test_struct_encoding_determinism (#683)
AratiGanesh
authored

Commits on Feb 9, 2026

Commits on Feb 10, 2026

Add `device test` unit tests for ROCm (JAX v0.9.0) (#705)
tsrw2048
authored

Commits on Feb 11, 2026

Commits on Feb 12, 2026

Update skip message version from 0.8.0 to 0.9.0 for test_is_finite on… (#702)
phambinhfin
authored

Commits on Mar 5, 2026

Fix HIP memory leaks in RNN kernels (#726)
magaonka-amd
authored