ROCm/jax

Pull Requests Commits

Fix export_back_compat_test: add rocsolver_potrf_ffi golden data

cj401-amd committed 43 days ago

4d0216fa

[ROCm] bypass hipSOLVER for Cholesky: call rocsolver potrf directly

cj401-amd committed 45 days ago

3cb46c7c

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

golechwierowicz committed 45 days ago

7b45bc03

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

Google-ML-Automation committed 45 days ago

b1e3e7b6

Reverts c3da028b6bd1e791c44e1d9d36b5771e1304cb46

junwhanahn committed 46 days ago

8e4f9e62

[bazel] visibility modifications

Jake VanderPlas committed 46 days ago

33a278b8

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

Google-ML-Automation committed 46 days ago

ba348c6c

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

yashk2810 committed 46 days ago

77944e22

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

Google-ML-Automation committed 46 days ago

862f2025

use empty not zeros in sum type representation

mattjj committed 46 days ago

15fca029

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

brills committed 46 days ago

f65eb481

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

yashk2810 committed 46 days ago

c3da028b

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

hawkinsp committed 46 days ago

46ea6d33

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

Google-ML-Automation committed 46 days ago

3e00b7cc

[pallas:mgpu] Fix flaky test_grid_dep_control

Google-ML-Automation committed 46 days ago

ecb0f7dd

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

hawkinsp committed 46 days ago

7e9ead37

Support tiling and swizzle transforms in GPU Interpret Mode.

jburnim committed 46 days ago

a4d60cd1

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

Google-ML-Automation committed 46 days ago

9172f4f2

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

Google-ML-Automation committed 46 days ago

b82f2ba9

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

Google-ML-Automation committed 46 days ago

5073cd4d

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

jburnim committed 46 days ago

f2e93613

Speedup check_sharding tests in export code.

hawkinsp committed 46 days ago

be69d848

[ROCm] CI improvements for ROCm

magaonka-amd committed 46 days ago

c178b348

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

bchetioui committed 46 days ago

b6e3b4fe

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

Google-ML-Automation committed 46 days ago

93c13aae

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

gnecula committed 46 days ago

63364b0f

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

bchetioui committed 46 days ago

2d1bd6e0

Adjust test tolerance for `testTakeAlongAxis` after compiler changes.

bchetioui committed 46 days ago

7f2b80f4

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

Google-ML-Automation committed 46 days ago

1274cf45

[bazel] visibility changes

Jake VanderPlas committed 47 days ago

bacac853

Older