flax

e48ce367 - Merge pull request #5455 from google:docs-examples--vit