Sergei Lebedev 5cef547eab Added support for lax.cond_p to Pallas Mosaic GPU lowering ...

PiperOrigin-RevId: 679156819

2024-09-26 08:20:53 -07:00

..

BUILD

Enable more H100 tests in CI.

2024-09-25 07:37:48 -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

Added support for lax.cond_p to Pallas Mosaic GPU lowering

2024-09-26 08:20:53 -07:00

ops_test.py

[Pallas] Add scalar f32 downcast test cases.

2024-09-25 11:25:59 -07:00

pallas_error_handling_test.py

[Pallas][XLA:Mosaic] Add python stack traces to Mosaic errors that occur in Pallas.

2024-08-12 14:42:48 -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