mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Disable the paged_attention test on TPU v5p.

Browse Source 
This test is failing in CI.

PiperOrigin-RevId: 694574616
This commit is contained in:
Peter Hawkins 2024-11-08 11:19:46 -08:00 committed by jax authors
parent 631bcd0942
commit 8f169e7fb5
1 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
tests/pallas/BUILD
View File

@ -390,6 +390,9 @@ jax_multiplatform_test(
jax_multiplatform_test(
name = "tpu_paged_attention_kernel_test",
srcs = ["tpu_paged_attention_kernel_test.py"],
disable_configs = [
"tpu_v5p_1x1",
],
enable_backends = ["tpu"],
shard_count = 5,
tags = [