flax

3711fa60 - Copybara import of the project:

Copybara import of the project: -- 36ee794e090e64e42574696647e303dd38fdafa2 by Cristian Garcia <cgarcia.e88@gmail.com>: [linen] fix DenseGeneral init PiperOrigin-RevId: 624873512