Peter Hawkins 5969e79908 Fix tests that ask for an accelerator but don't use it. ...

* Delete custom_object_test, since it is disabled and has been ever since jax.Array was introduced in JAX 0.4.0.
* custom_linear_solve_test was over-sharded, leading to some shards not having any test cases. Even unsharded it completes in under 65s on every platform we have.
* config_test and pallas splash attention mask test only tested helpers and didn't need a TPU.

PiperOrigin-RevId: 679711664

2024-09-27 13:36:23 -07:00

..

BUILD

Fix tests that ask for an accelerator but don't use it.

2024-09-27 13:36:23 -07:00

export_back_compat_pallas_test.py

[Pallas] Update export compatibility tests

2024-09-24 09:17:56 -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

[Pallas] Fully skip GPU attention tests on win32.

2024-09-09 10:21:33 -07:00

gpu_ops_test.py

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

indexing_test.py

[Pallas] Fix array indexing error when dimension size is not a multiple of stride

2024-09-10 02:52:55 +01:00

mosaic_gpu_test.py

[Pallas/MGPU] Implement block spec evaluation correctly

2024-09-26 09:15:12 -07:00

ops_test.py

[Pallas:TPU] Fix lowering of convert_element_type(int32) -> bool.

2024-09-27 02:15:35 -07: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

Reverts 66a3f87a24016594794c2ee289826baed5e979a4

2024-08-20 15:13:09 -07:00

pallas_shape_poly_test.py

Reverts 0d058ce86f04a44a51abba1261768fb46edf69d9

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

pallas_test.py

[pallas] Add support for unblocked mode (without padding) in Triton lowering.

2024-09-23 11:21:54 -07:00

pallas_vmap_test.py

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -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 TPU] Refactor ref indexers to transforms and support ref bitcast.

2024-09-15 17:53:29 -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_async_test.py

[Pallas TPU] Add support for passing in and returning semaphores

2024-09-18 20:39:43 -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

Temporarily disable broken tests in tpu_pallas_pipeline_test.py

2024-08-08 14:04:04 -07:00

tpu_pallas_random_test.py

[pallas] Renamed platform-specific tests.

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

tpu_pallas_test.py

[pallas] Add support for unblocked mode (without padding) in Triton lowering.

2024-09-23 11:21:54 -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