flax

add promote_dtype to all standard layers

#5080

Merged