Why is ada_lin computation (e.g., cond_BD) and get_logits explicitly required to be in float32?

[Deer1111]

Hello author

I was looking through the code and noticed that in the ada_lin module, some calculations (like cond_BD) are explicitly flagged to use float32. There's even a comment that says:

"cond_BD should be float32"

I'm curious — what's the main reason for requiring float32 here? What would go wrong if we used half-precision (like float16 or bfloat16) instead?

My guess is it's probably about numerical stability… maybe half-precision could cause overflow/underflow in these ops? Would love to hear the reasoning behind it!

Thanks in advance for the insight! 🙌