vision_transformer

Merge pull request #266 from limesqueezy:main

Modified `vit_jax.ipynb` as requested.

Update vit_jax.ipynb

Internal change

Factor out the "positional embedding interpolation" logic into a standalone public function.