Try launching unit tests on TPUs from CI

[dlwh]

No description provided.

[dlwh]

cc @rjpower since I fixed up a few tests on TPU (mostly just scaled down some initializers?)

[rjpower]

Curious... not surprising that the numerics would be slightly different for splash attention but interesting how changing the initialization helps that much. I think we expect a lot of noise through the matmul unit so this is likely what we expect.

I wonder if the scaling is implicitly adjusting the tolerance here, if that makes sense? That is, in the original, our output might be on the scale of 1000.00 {+- 0.5}, and now it's 1.00000 { +- 0.5.}? We're accurate to the same number of digits in both cases, just the digits that we're testing against are different. As a dumb example let's assume one attention went through bf16 and the other didn't, we'd see something like this through the matmuls:

import numpy as np
from paddle_bfloat import bfloat16

x = np.random.normal(-1, 1, 1000)
xs = x * 0.02

b = x.astype(bfloat16)
bs = xs.astype(bfloat16)

print(x@x - b@b)
print(xs@xs - bs@bs)

-0.1878888127475875
8.110935221994353e-05

[dlwh]

one of the fixed tests is actually my "pure jax" flash attention vs vanilla dot product, so it's the same algorithm on both CPU and TPU, but the matmuls are different enough on TPU that it matters, even with HIGHEST.

I mostly agree with your hypothesis, though I framed it in terms of "floating point numbers are more precise near 0"

Floating point is so annoying