flax

9c419680 - Don't pass linen_meta_type argument when creating AxisMetadata subclasses.