George Necula 3e5e947542 Move some backwards compatibility tests from jax_triton to jax/pallas. ...

While doing this I moved `matmul.py` to `jax/experimental/pallas/ops/tpu`

PiperOrigin-RevId: 660341331

2024-08-07 05:00:29 -07:00

..

BUILD

Move some backwards compatibility tests from jax_triton to jax/pallas.

2024-08-07 05:00:29 -07:00

export_back_compat_pallas_test.py

Move some backwards compatibility tests from jax_triton to jax/pallas.

2024-08-07 05:00:29 -07:00

export_pallas_test.py

[pallas] Add lowering errors for block shapes that are not supported.

2024-07-29 06:49:27 -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

Handle bool comparisons.

2024-08-06 05:37:35 -07:00

pallas_shape_poly_test.py

Reverts 0d058ce86f04a44a51abba1261768fb46edf69d9

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

pallas_test.py

[pallas] Improve error and debugging messages with source locations

2024-08-05 04:23:55 -07:00

pallas_vmap_test.py

[pallas] Improve error and debugging messages with source locations

2024-08-05 04:23:55 -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_paged_attention_kernel_test.py

Skip some Pallas tests that fail on TPUv6.

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

tpu_pallas_distributed_test.py

[Pallas] Correctly handle asymmetrical remote DMA dst_ref indexing in interpret mode.

2024-07-31 09:10:53 -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

Move CostEstimate from pltu to pl

2024-08-05 08:18:01 -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