ROCm/jax

Pull Requests Commits

[Mosaic GPU] Add swizzle in `SMEMTransforms`.

golechwierowicz committed 46 days ago

48f94e65

[Mosaic GPU] Add support for indexing untiled dimensions to `FragmentedArray`.

allanrenucci committed 46 days ago

6302b470

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

Google-ML-Automation committed 46 days ago

6afcac46

[Mosaic TPU] Remove some unnecessary skips on packed type reshape tests.

yueshengys committed 46 days ago

9bd8cb58

Support prefetched inputs in Pallas pipelines

Google-ML-Automation committed 46 days ago

4d4536de

Merge pull request #38212 from Intel-tensorflow:oneapi_test_utils

Google-ML-Automation committed 46 days ago

8befb14c

Add cudnn sdpa flex attention in JAX SDPA API

Cjkkkk committed 46 days ago

53c96c17

Skip a test that is setting a flag only just added to XLA.

hawkinsp committed 46 days ago

be255b0e

Refactor Windows Docker GCE metadata access

nitins17 committed 46 days ago

6beff407

Make persistent compilation cache support host callbacks

skye-ant committed 46 days ago

Verified 54156b8c

Elevating the Pallas docs to higher visibility.

IvyZX committed 46 days ago

8c196a35

Remove an unnecessary tpu version skip, now that libtpu 0.0.42 is the minimum.

hawkinsp committed 46 days ago

33aaddf0

Remove deprecated is_cloud_tpu_at_least version guards.

hawkinsp committed 46 days ago

9ddd12f8

Merge pull request #38480 from hawkinsp:libtpu_pin

Google-ML-Automation committed 46 days ago

90e60ab5

Merge pull request #38446 from hawkinsp:portserver

Google-ML-Automation committed 46 days ago

e518a633

Bump libtpu pin to 0.0.42.

hawkinsp committed 46 days ago

f4912996

Merge pull request #38478 from magaonka-amd:magaonka/skip-conv-batching-miopen

Google-ML-Automation committed 46 days ago

ccb76bb2

add test utilities functions for oneapi

kranipa committed 46 days ago

00b8c5db

Merge pull request #38452 from jax-ml:push-zyumqxqoktly

Google-ML-Automation committed 46 days ago

d69e64ef

[ROCm] Skip testConvGeneralDilated due to XLA/MIOpen boundary bug

magaonka-amd committed 46 days ago

a49224b6

Merge pull request #38407 from hawkinsp:libtpu

Google-ML-Automation committed 47 days ago

e5782eb4

Improve sharding IndivisibleError error message by adding pytree key paths and source location info.

yashk2810 committed 47 days ago

57ba662a

Add a jtu.is_libtpu_at_least version guard.

hawkinsp committed 47 days ago

24b57975

Add a workaround for linux aarch64 python 3.14t pytest builds.

hawkinsp committed 47 days ago

2951d463

[Pallas:MGPU] Add support for batch dimensions to `async_store_tmem`.

allanrenucci committed 47 days ago

7db3a70e

[Mosaic:GPU] Disable set_xla_gpu_enable_allocator_spatial_partitioning for nvshmem test case, since nvshmem can't be used within the same allocation.

PatriosTheGreat committed 47 days ago

5597e903

Add functionality to blanket disable the custom call safety checks in export.

WillFroom committed 47 days ago

728e4bfe

Merge pull request #38432 from vfdev-5:pallas-emit-pipeline-raises-error-empty-grid

Google-ML-Automation committed 47 days ago

26113fa5

Pallas emit_pipeline raises ValueError on empty grid

vfdev-5 committed 47 days ago

c15bacc8

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

Google-ML-Automation committed 47 days ago

31a7aaba

Newer Older