George Necula 6d53aaf7d0 [pallas] Improve the error localization ...

* Add the source location information for the index map function to
    `BlockMapping`.
  * Removed the `compute_index` wrapper around the index_map, so that
    we can get the location information for the index_map, not the wrapper.
  * Added source location to the errors related to index map functions.
  * Added an error if the index map returns something other than integer
    scalars.
  * Construct BlockSpec origins for arguments using JAX helper functions
    to get argument names
  * Removed redundant API error tests from tpu_pallas_test.py

2024-07-30 14:11:57 +02:00

..

BUILD

[pallas:mosaic_gpu] Add support for run_scoped

2024-07-23 16:13:00 -07:00

export_back_compat_pallas_test.py

Migrated `pl.BlockSpec` uses in JAX to the new argument order

2024-07-02 05:16:54 -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

Merge pull request #22716 from superbobry:pallas

2024-07-29 14:50:01 -07:00

pallas_shape_poly_test.py

Reverts 0d058ce86f04a44a51abba1261768fb46edf69d9

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

pallas_test.py

[pallas] Improve the error localization

2024-07-30 14:11:57 +02:00

pallas_vmap_test.py

Reverts 0d058ce86f04a44a51abba1261768fb46edf69d9

2024-07-25 01:50:36 -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_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

[pallas] Improve the error localization

2024-07-30 14:11:57 +02: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