flax

c2865a17 - Minor improvements in dtype handling, broadcast option for dropout.

Minor improvements in dtype handling, broadcast option for dropout. We add a dtype specification to Embed layer, make Adafactor use float32 state consistently, and add a broadcasting option to the Dropout layer. PiperOrigin-RevId: 337420935