ROCm/jax

Pull Requests Commits

devtest: continuous CI for 1-GPU ROCm pytest off fork branches

magaonka-amd committed 1 day ago

1a4abdc2

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

Google-ML-Automation committed 1 day ago

9172f4f2

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

Google-ML-Automation committed 1 day ago

b82f2ba9

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

Google-ML-Automation committed 1 day ago

5073cd4d

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

jburnim committed 1 day ago

f2e93613

Speedup check_sharding tests in export code.

hawkinsp committed 1 day ago

be69d848

[ROCm] CI improvements for ROCm

magaonka-amd committed 1 day ago

c178b348

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

bchetioui committed 2 days ago

b6e3b4fe

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

Google-ML-Automation committed 2 days ago

93c13aae

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

gnecula committed 2 days ago

63364b0f

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

bchetioui committed 2 days ago

2d1bd6e0

Adjust test tolerance for `testTakeAlongAxis` after compiler changes.

bchetioui committed 2 days ago

7f2b80f4

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

Google-ML-Automation committed 2 days ago

1274cf45

[bazel] visibility changes

Jake VanderPlas committed 2 days ago

bacac853

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

Google-ML-Automation committed 3 days ago

9863d3af

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

jburnim committed 3 days ago

a5634cb9

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

superbobry committed 4 days ago

0e23efd6

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

Google-ML-Automation committed 4 days ago

954cb316

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

Google-ML-Automation committed 4 days ago

f1e823f1

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

Google-ML-Automation committed 4 days ago

fc9d6bfd

Add regression test for jnp.corrcoef NumPy 2.2 behavior

adityasingh2400 committed 4 days ago

Verified 33e47b56

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

jburnim committed 4 days ago

950a9c7f

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

Google-ML-Automation committed 4 days ago

f1ff9967

Remove outdated ROCm skip that depends on ROCm versions < 6.4 (#776)

gulsumgudukbay committed 4 days ago

13e3954a

[Pallas] Set preferred_element_type.

WindQAQ committed 4 days ago

8275f781

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

Google-ML-Automation committed 4 days ago

f078f550

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

superbobry committed 4 days ago

f3b27574

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

Google-ML-Automation committed 4 days ago

e9fdb286

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

hawkinsp committed 4 days ago

1587193c

[Mosaic GPU] Add support for PDL primitives to Pallas Mosaic GPU

rajaanand-g committed 4 days ago

702e27b2

Older