ROCm/jax

Pull Requests Commits

[ROCm] rocsolver_potrf_ffi: dedicated rocblas_handle pool for potrf

cj401-amd committed 3 hours ago

7370ae44

[ROCm] Add rocsolver_potrf_ffi: bypass hipSOLVER for Cholesky on ROCm

cj401-amd committed 3 hours ago

c26560f2

[Mosaic GPU] Safely construct certain layouts instead of failing fast with exception.

golechwierowicz committed 4 hours ago

7b45bc03

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

Google-ML-Automation committed 5 hours ago

b1e3e7b6

Reverts c3da028b6bd1e791c44e1d9d36b5771e1304cb46

junwhanahn committed 12 hours ago

8e4f9e62

[bazel] visibility modifications

Jake VanderPlas committed 12 hours ago

33a278b8

Merge pull request #38279 from mattjj:empty-not-zeros-in-sum-type-representation

Google-ML-Automation committed 13 hours ago

ba348c6c

Improve error message for sharding rules when the dimension does not divide mesh axes it's sharded on

yashk2810 committed 13 hours ago

77944e22

Merge pull request #38168 from mattjj:accum-ct-avals

Google-ML-Automation committed 13 hours ago

862f2025

use empty not zeros in sum type representation

mattjj committed 14 hours ago

15fca029

Bug fix: do not create real (non-symbolic) zero arrays for tangents.

brills committed 14 hours ago

f65eb481

Use `lax.empty` in `_empty_array` and add `out_memory_space` to `lax.empty`.

yashk2810 committed 14 hours ago

c3da028b

[pallas] When unrolling loops during mosaic lowering, emit the loop body out of line and repeatedly inline it.

hawkinsp committed 15 hours ago

46ea6d33

[Pallas:MGPU] Allow dynamic shapes in nd_loop and planar_snake.

Google-ML-Automation committed 15 hours ago

3e00b7cc

[pallas:mgpu] Fix flaky test_grid_dep_control

Google-ML-Automation committed 15 hours ago

ecb0f7dd

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

hawkinsp committed 15 hours ago

7e9ead37

Support tiling and swizzle transforms in GPU Interpret Mode.

jburnim committed 16 hours ago

a4d60cd1

Merge pull request #38200 from ROCm:remove-outdated-rocm-skip-upstream

Google-ML-Automation committed 16 hours ago

9172f4f2

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

Google-ML-Automation committed 19 hours ago

b82f2ba9

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

Google-ML-Automation committed 20 hours ago

5073cd4d

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

jburnim committed 21 hours ago

f2e93613

Speedup check_sharding tests in export code.

hawkinsp committed 23 hours ago

be69d848

[ROCm] CI improvements for ROCm

magaonka-amd committed 23 hours ago

c178b348

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

bchetioui committed 1 day ago

b6e3b4fe

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

Google-ML-Automation committed 1 day ago

93c13aae

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

gnecula committed 1 day ago

63364b0f

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

bchetioui committed 1 day ago

2d1bd6e0

Adjust test tolerance for `testTakeAlongAxis` after compiler changes.

bchetioui committed 1 day ago

7f2b80f4

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

Google-ML-Automation committed 1 day ago

1274cf45

[bazel] visibility changes

Jake VanderPlas committed 1 day ago

bacac853

Older