ROCm/jax

Pull Requests Commits

[ROCm] Load rocsolver 3.33.0 via dlopen for Cholesky potrf

cj401-amd committed 30 days ago

a5ed5e0f

jaxlib/gpu: query rocsolver potrf workspace with correct element type

cj401-amd committed 30 days ago

fd9396e6

jaxlib/gpu: pre-allocate rocBLAS workspace for rocsolver potrf on ROCm

cj401-amd committed 30 days ago

3663a752

[ROCm] Use native rocsolver potrf for Cholesky instead of hipSOLVER

cj401-amd committed 30 days ago

27d5c6bf

Merge pull request #38170 from Intel-tensorflow:ashahba/oneapi-ci-build

Google-ML-Automation committed 30 days ago

b82f2ba9

Merge pull request #38258 from magaonka-amd:ci/rocm-default-allocator

Google-ML-Automation committed 30 days ago

5073cd4d

Syle fixes to https://github.com/jax-ml/jax/pull/38227 .

jburnim committed 30 days ago

f2e93613

Speedup check_sharding tests in export code.

hawkinsp committed 30 days ago

be69d848

[ROCm] CI improvements for ROCm

magaonka-amd committed 30 days ago

c178b348

[Mosaic GPU] Warp-level MMA no longer takes `b` as transposed.

bchetioui committed 30 days ago

b6e3b4fe

[pallas:mgpu] Add jax.block_until_ready to ensure PTX is captured in test_griddepcontrol

Google-ML-Automation committed 30 days ago

93c13aae

PR #37220: [debug_info] Fix result_paths usage in _check_returned_jaxtypes

gnecula committed 30 days ago

63364b0f

Roll back PR #36850. `DotProductAttentionTest.test_sdpa_flex_attention` returns

bchetioui committed 30 days ago

2d1bd6e0

Adjust test tolerance for `testTakeAlongAxis` after compiler changes.

bchetioui committed 31 days ago

7f2b80f4

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

Google-ML-Automation committed 31 days ago

1274cf45

[bazel] visibility changes

Jake VanderPlas committed 31 days ago

bacac853

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

Google-ML-Automation committed 32 days ago

9863d3af

[Pallas] In GPU Interpret Mode, replace ordered IO callbacks with unordered

jburnim committed 32 days ago

a5634cb9

[state] Rolling forward `discharge_state2`->`discharge_state`

superbobry committed 32 days ago

0e23efd6

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

Google-ML-Automation committed 33 days ago

954cb316

In Primitive.bind, canonicalize first before calling typeof and fuse the loops that compute them.

Google-ML-Automation committed 33 days ago

f1e823f1

Merge pull request #37966 from adityasingh2400:adityasingh2400/corrcoef-numpy-22-parity

Google-ML-Automation committed 33 days ago

fc9d6bfd

Add regression test for jnp.corrcoef NumPy 2.2 behavior

adityasingh2400 committed 33 days ago

Verified 33e47b56

[Pallas] In TPU Interpret Mode, notify Semaphore.cv when a task is queued.

jburnim committed 33 days ago

950a9c7f

Merge pull request #38204 from Sumu004:fix/array-platform-type-stub

Google-ML-Automation committed 33 days ago

f1ff9967

[Pallas] Set preferred_element_type.

WindQAQ committed 33 days ago

8275f781

[state] Use `discharge_state2` everywhere and drop the 2

Google-ML-Automation committed 33 days ago

f078f550

[state] Use `discharge_state2` everywhere and drop the 2

superbobry committed 33 days ago

f3b27574

Merge pull request #38208 from ROCm:use-rocm-tag-prefix-s3

Google-ML-Automation committed 33 days ago

e9fdb286

In Primitive.bind, canonicalize first before calling typeof and fuse the loops that compute them.

hawkinsp committed 33 days ago

1587193c

Older