unit test for tpu_attention_test.py fails

[lkolluru05]

output_0.5.3.txt

Hi,

I ran this test https://github.com/apple/axlearn/blob/main/axlearn/common/flash_attention/tpu_attention_test.py on v6e-4 machine and I got 4 tests failing

=========================== short test summary info ============================ SKIPPED [144] axlearn/common/flash_attention/test_utils.py:47: segment ids require kv_seq_len == q_seq_len =========== 4 failed, 300 passed, 144 skipped in 1366.78s (0:22:46) ============

Below I am attaching the test results

output_0.4.38.txt