Sergei Lebedev 99a8b92a7f Fixed Pallas Mosaic GPU tests ...

* Migrated to the new barrier APIs
* Fixed scratch view casting logic, it previously didn't work for >1 view

PiperOrigin-RevId: 655937541

2024-07-25 06:51:13 -07:00

..

BUILD

[pallas:mosaic_gpu] Add support for run_scoped

2024-07-23 16:13:00 -07:00

export_back_compat_pallas_test.py

Migrated `pl.BlockSpec` uses in JAX to the new argument order

2024-07-02 05:16:54 -07:00

export_pallas_test.py

[pallas] Move the import of tpu_custom_call from tests to a more central place.

2024-07-08 04:27:41 -07:00

gpu_attention_test.py

Skip pallas/gpu_attention_test.py on TPU.

2024-07-24 08:24:57 -07:00

gpu_ops_test.py

Reverts 0d058ce86f04a44a51abba1261768fb46edf69d9

2024-07-25 01:50:36 -07:00

indexing_test.py

Add some hypothesis testing utilities and developer documentation.

2024-07-15 17:05:32 +02:00

mosaic_gpu_test.py

Fixed Pallas Mosaic GPU tests

2024-07-25 06:51:13 -07:00

ops_test.py

[Pallas] Add better reduction support.

2024-07-09 11:03:17 -07:00

pallas_shape_poly_test.py

Reverts 0d058ce86f04a44a51abba1261768fb46edf69d9

2024-07-25 01:50:36 -07:00

pallas_test.py

Reverts 0d058ce86f04a44a51abba1261768fb46edf69d9

2024-07-25 01:50:36 -07:00

pallas_vmap_test.py

Reverts 0d058ce86f04a44a51abba1261768fb46edf69d9

2024-07-25 01:50:36 -07:00

tpu_all_gather_test.py

[pallas] Renamed platform-specific tests.

2024-07-10 08:23:06 -07:00

tpu_gmm_test.py

[pallas] Renamed platform-specific tests.

2024-07-10 08:23:06 -07:00

tpu_ops_test.py

[pallas] Renamed platform-specific tests.

2024-07-10 08:23:06 -07:00

tpu_paged_attention_kernel_test.py

Skip some Pallas tests that fail on TPUv6.

2024-07-23 07:16:24 -07:00

tpu_pallas_mesh_test.py

[pallas] Use the same primitive run_scoped_p for moth mosaic and mosaic_gpu

2024-07-24 17:14:30 -07:00

tpu_pallas_pipeline_test.py

[pallas] Use the same primitive run_scoped_p for moth mosaic and mosaic_gpu

2024-07-24 17:14:30 -07:00

tpu_pallas_random_test.py

[pallas] Renamed platform-specific tests.

2024-07-10 08:23:06 -07:00

tpu_pallas_test.py

Reverts 0d058ce86f04a44a51abba1261768fb46edf69d9

2024-07-25 01:50:36 -07:00

tpu_splash_attention_kernel_test.py

[Pallas/TPU] Fix bug with LocalMask grid shrinking

2024-07-25 04:05:57 -07:00

tpu_splash_attention_mask_test.py

[pallas] Renamed platform-specific tests.

2024-07-10 08:23:06 -07:00