py: Write P_expected, don't rewrite vars

2024-09-04 23:35:52 -07:00
parent dcd69ea304
commit bde6f0ea2e
2 changed files with 9 additions and 8 deletions
--- a/tests/kernel/tensor/flash_attn.py
+++ b/tests/kernel/tensor/flash_attn.py
@@ -133,6 +133,7 @@ if __name__ == "__main__":
            print('P_expected:')
            print(P)
            P.astype('float32').tofile("P_expected.bin")
+            P.transpose([1, 0]).astype('float32').tofile("P_expected.col.bin")

        rowsum_this = np.sum(P, axis=1)
        x = rowmax_prev - rowmax_this