flax

dc218b40 - Rename abstract_with_sharding to as_abstract