Adam Paszke d2a5e8d072 [Mosaic TPU] Add support for integer truncation from packed types ...

PiperOrigin-RevId: 714048232

2025-01-10 07:40:55 -08:00

..

BUILD

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

2025-01-10 04:35:38 -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

Save residuals in the decode attention pallas kernel

2024-12-02 15:09:16 +00:00

gpu_ops_test.py

Adding more tests for multi-head attention

2024-12-09 20:49:06 +00: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

[pallas:mosaic_gpu] Tests now pass with x64 enabled

2025-01-10 04:48:36 -08:00

ops_test.py

[Mosaic TPU] Add support for integer truncation from packed types

2025-01-10 07:40:55 -08: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

Reverts dbe9ccd6dccd83c365021677c7e17e843d4559c4

2025-01-10 03:40:18 -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] Deprecate dictionary compiler_params in favor of dataclass.

2024-11-21 23:34:32 -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

[Mosaic:TPU] Vreg-slice-aligned offset changes with scratch retiling

2024-12-23 13:05:14 -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