jax

872e6c0e - Merge pull request #25766 from carlosgmartin:nn_initializers_variance_scaling_mode_fan_geo_avg