flax

719f6879 - Update gspmd guide to mention Variable init with eager_sharding