Make the default matmul precision float32 even on TPUs

[shoyer]

Follow-up from #2161:

The default low precision is a bit of a footgun, at least when doing anything that isn't implementing a neural net layer. In my opinion, it would be much safer to use "highest" precision by default (which isn't that much slower) on float32 data. Neural net libraries, of course, can default to lower precision, so this really only effects users who directly use NumPy APIs or the @ infix operator.

[j-towns]

I would argue this is also a footgun for some neural net use cases. Both of the two times that I’ve carefully re-implemented a model in JAX, I’ve found that performance is worse than expected with the default precision and it took me some time to realise that precision was the cause.

As a (temporary) improvement on the current situation, we could at least add some information on this issue to point 5 of the ‘Current Gotchas’ in the readme.

[shoyer]

#6143 added a config flag, so this should in principle be easier to change now.

[j-towns]

As @juliuskunze just pointed out to me, this is the only difference in semantics (that we can think of) between the GPU, CPU and TPU backends. 'Backend transparency' is a valuable property in mine and Julius's opinion, and given how close JAX is to achieving it (if this really is the only case where semantics differ significantly), it's surely worth changing the TPU default for all matmul-like ops (including conv) to closely approximate GPU and CPU behaviour. I understand this will harm 'default' speed, but I think we can mitigate that by making clear the availability of the bfloat16 option.

[jonbarron]

My team shot itself in the foot last week for the ~fourth time due to matmult defaulting to bfloat16. This issue continues to be my biggest/only grievance with JAX.

[mattjj]

Just got another +1 for this issue. Another foot lost.

[patrickvonplaten]

+1 for this issue from the Transformers team for one of the most popular architectures for speech recognition (Wav2Vec2) - see: huggingface/transformers#15754

[inoryy]

+1 to setting default to f32.

Having had the pleasure of shooting both my feet with either performance gotchas and numerics gotchas, I'd wholeheartedly prefer debugging the performance ones as they're much easier to spot.