Workaround https://github.com/google/jax/issues/5406 in DotGeneral with masking by casting to float
Workaround https://github.com/google/jax/issues/5406 in DotGeneral with masking by casting to float

PiperOrigin-RevId: 355533576

google/neural-tangents