flax

dc841b04 - Fix dtype handling for Adam and LAMB optimizers in 64bit mode.

Fix dtype handling for Adam and LAMB optimizers in 64bit mode. This prevents accidental upcasting of parameters due to wonky JAX integer + floating-point py-scalar + array dtype rules.