Ruturaj Vaidya 790fe78b82

Skip tests on ROCm (#235)

2025-02-20 09:41:11 -06:00

..

BUILD

[Mosaic GPU] Enable x64 tests for mosaic gpu.

2025-01-14 13:02:48 -08:00

export_back_compat_pallas_test.py

[Mosaic TPU] Guard tests for new features by the libtpu version

2024-12-19 05:04:09 -08:00

export_pallas_test.py

[Mosaic TPU] Guard tests for new features by the libtpu version

2024-12-19 05:04:09 -08:00

gpu_attention_test.py

[Pallas] Fix GQA triton kernel test.

2025-01-14 16:40:55 -08:00

gpu_ops_test.py

Skip tests on ROCm (#235)

2025-02-20 09:41:11 -06:00

indexing_test.py

Always suppress the differing_executors Hypothesis health check

2024-12-19 08:00:53 -08:00

mgpu_attention_test.py

[Pallas:MGPU] Add tests for attention with non-trivial batch size

2024-12-03 03:58:15 -08:00

mosaic_gpu_test.py

[Mosaic GPU] Enable x64 tests for mosaic gpu.

2025-01-14 13:02:48 -08:00

ops_test.py

Skip tests on ROCm (#235)

2025-02-20 09:41:11 -06:00

pallas_cost_estimate_test.py

[Pallas] Add support for run_state to cost estimator.

2024-12-06 10:36:02 -08:00

pallas_error_handling_test.py

[XLA:Mosaic][Pallas] Enable vector.ExtractOp for non-zero indices.

2024-09-26 13:57:45 -07:00

pallas_jumble_test.py

[pallas] fix jumble test flakiness

2024-12-05 15:31:23 -08:00

pallas_shape_poly_test.py

Reverts 0d058ce86f04a44a51abba1261768fb46edf69d9

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

pallas_test.py

[pallas][mosaic kernel export] Add initial support for exporting a dynamic shapes (placeholder bound) kernel out of mosaic, via pallas as both MLIR and jaxpr.

2025-01-14 20:34:11 -08:00

pallas_vmap_test.py

[pallas] Do not skip vmap tests on GPU when x64 is enabled

2024-11-20 05:08:23 -08:00

tpu_all_gather_test.py

[pallas] Renamed platform-specific tests.

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

tpu_gmm_test.py

Make gmm TPU kernel tests significantly cheaper

2024-12-17 08:38:59 -08:00

tpu_ops_test.py

[Pallas] Improve testing for lowering of dtype conversions + fix uncovered bugs

2025-01-10 04:35:38 -08: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_async_test.py

[Pallas] Add state discharge rule for pallas_call

2024-10-01 16:30:56 -07:00

tpu_pallas_distributed_test.py

[Mosaic TPU] Add support for modeling loads/stores and fix minor issues in model extraction

2024-12-05 07:06:19 -08:00

tpu_pallas_pipeline_test.py

[Pallas TPU] Temporarily strengthen restrictions on Pallas tests

2025-01-15 02:32:14 -08:00

tpu_pallas_random_test.py

Add Pallas Philox implementation.

2024-12-17 15:37:08 -08:00

tpu_pallas_state_test.py

[Pallas] Use core_map instead of shard_map for Shmallas

2024-10-15 03:26:58 -07:00

tpu_pallas_test.py

Jax: Stop returning a list of cost-analyses.

2025-01-15 09:53:59 -08:00

tpu_splash_attention_kernel_test.py

Reduce the max_examples for splash attention tests

2024-12-17 09:06:19 -08:00

tpu_splash_attention_mask_test.py

Use boolean values for partial mask blocks in the splash attention kernel.

2024-12-11 14:59:30 -08:00