flax

2d32f287 - replaced zeros and ones initializers with zeros_init and ones_init builder functions