flax

187b910e - Copybara import of the project:

Copybara import of the project: -- 525c3d87eff8ad1c4e4fe507464e5bb251cb69e0 by Cristian Garcia <cgarcia.e88@gmail.com>: add promote_dtype as a config option for multiple layers PiperOrigin-RevId: 735924369