Flax implementation of Vit Lite?

[dhyani15]

Hi team,

Thanks for open sourcing your work, may I ask if you might know any implementation of these models in Flax/Jax?

Best,
Naren

[stevenwalton]

ViT-Lite is simply the normal ViT with reduced layers and smaller patch sizes. See Section 3.2. Any existing Flax/Jax ViT model can be trivially adjusted to fit this.

As far as CCT, if you find or implement it (not a significantly difficult task) then we'd be happy to add it to our README in addition to the other third party implementations.