vibe out a vit with sparse gating with a thresholded l1 loss for personal use

* added vit-5 implementation

* updated images for readme for vit-5

.github/FUNDING.yml

@@ -1,3 +0,0 @@
-# These are supported funding model platforms
-
-github: [lucidrains]

.github/workflows/python-publish.yml

@@ -1,31 +0,0 @@
-# This workflows will upload a Python Package using Twine when a release is created
-# For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries
-
-name: Upload Python Package
-
-on:
-  release:
-    types: [created]
-
-jobs:
-  deploy:
-
-    runs-on: ubuntu-latest
-
-    steps:
-    - uses: actions/checkout@v2
-    - name: Set up Python
-      uses: actions/setup-python@v2
-      with:
-        python-version: '3.x'
-    - name: Install dependencies
-      run: |
-        python -m pip install --upgrade pip
-        pip install setuptools wheel twine
-    - name: Build and publish
-      env:
-        TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }}
-        TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }}
-      run: |
-        python setup.py sdist bdist_wheel
-        twine upload dist/*

.github/workflows/python-test.yml

@@ -1,34 +0,0 @@
-# This workflow will install Python dependencies, run tests and lint with a variety of Python versions
-# For more information see: https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions
-
-name: Test
-
-on:
-  push:
-    branches: [ main ]
-  pull_request:
-    branches: [ main ]
-
-jobs:
-  build:
-
-    runs-on: ubuntu-latest
-    strategy:
-      matrix:
-        python-version: [3.7, 3.8, 3.9]
-
-    steps:
-    - uses: actions/checkout@v2
-    - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v2
-      with:
-        python-version: ${{ matrix.python-version }}
-    - name: Install dependencies
-      run: |
-        python -m pip install --upgrade pip
-        python -m pip install pytest
-        python -m pip install wheel
-        if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
-    - name: Test with pytest
-      run: |
-        python setup.py test

.gitignore

@@ -127,3 +127,6 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+
+# scripts
+*.sh

README.md

@@ -25,6 +25,7 @@
 - [MaxViT](#maxvit)
 - [NesT](#nest)
 - [MobileViT](#mobilevit)
+- [XCiT](#xcit)
 - [Masked Autoencoder](#masked-autoencoder)
 - [Simple Masked Image Modeling](#simple-masked-image-modeling)
 - [Masked Patch Prediction](#masked-patch-prediction)
@@ -48,7 +49,7 @@
 
 ## Vision Transformer - Pytorch
 
-Implementation of <a href="https://openreview.net/pdf?id=YicbFdNTTy">Vision Transformer</a>, a simple way to achieve SOTA in vision classification with only a single transformer encoder, in Pytorch. Significance is further explained in <a href="https://www.youtube.com/watch?v=TrdevFK_am4">Yannic Kilcher's</a> video. There's really not much to code here, but may as well lay it out for everyone so we expedite the attention revolution.
+Implementation of <a href="https://openreview.net/pdf?id=YicbFdNTTy">Vision Transformer</a>, a simple way to achieve SOTA in vision classification with only a single transformer encoder, in Pytorch. Significance is further explained in <a href="https://www.youtube.com/watch?v=TrdevFK_am4">Yannic Kilcher's</a> video. There's really not much to code here, but may as well lay it out for everyone so we expedite the [attention](https://www.youtube.com/watch?v=eMlx5fFNoYc) revolution.
 
 For a Pytorch implementation with pretrained models, please see Ross Wightman's repository <a href="https://github.com/rwightman/pytorch-image-models">here</a>.
 
@@ -89,26 +90,26 @@ preds = v(img) # (1, 1000)
 
 ## Parameters
 
-- `image_size`: int.  
+- `image_size`: int.
 Image size. If you have rectangular images, make sure your image size is the maximum of the width and height
-- `patch_size`: int.  
-Number of patches. `image_size` must be divisible by `patch_size`.  
+- `patch_size`: int.
+Size of patches. `image_size` must be divisible by `patch_size`.
 The number of patches is: ` n = (image_size // patch_size) ** 2` and `n` **must be greater than 16**.
-- `num_classes`: int.  
+- `num_classes`: int.
 Number of classes to classify.
-- `dim`: int.  
+- `dim`: int.
 Last dimension of output tensor after linear transformation `nn.Linear(..., dim)`.
-- `depth`: int.  
+- `depth`: int.
 Number of Transformer blocks.
-- `heads`: int.  
-Number of heads in Multi-head Attention layer. 
-- `mlp_dim`: int.  
-Dimension of the MLP (FeedForward) layer. 
-- `channels`: int, default `3`.  
-Number of image's channels. 
-- `dropout`: float between `[0, 1]`, default `0.`.  
-Dropout rate. 
-- `emb_dropout`: float between `[0, 1]`, default `0`.  
+- `heads`: int.
+Number of heads in Multi-head Attention layer.
+- `mlp_dim`: int.
+Dimension of the MLP (FeedForward) layer.
+- `channels`: int, default `3`.
+Number of image's channels.
+- `dropout`: float between `[0, 1]`, default `0.`.
+Dropout rate.
+- `emb_dropout`: float between `[0, 1]`, default `0`.
 Embedding dropout rate.
 - `pool`: string, either `cls` token pooling or `mean` pooling
 
@@ -179,6 +180,56 @@ preds = v(images) # (5, 1000) - 5, because 5 images of different resolution abov
 
 ```
 
+Or if you would rather that the framework auto group the images into variable lengthed sequences that do not exceed a certain max length
+
+```python
+images = [
+    torch.randn(3, 256, 256),
+    torch.randn(3, 128, 128),
+    torch.randn(3, 128, 256),
+    torch.randn(3, 256, 128),
+    torch.randn(3, 64, 256)
+]
+
+preds = v(
+    images,
+    group_images = True,
+    group_max_seq_len = 64
+) # (5, 1000)
+```
+
+Finally, if you would like to make use of a flavor of NaViT using <a href="https://pytorch.org/tutorials/prototype/nestedtensor.html">nested tensors</a> (which will omit a lot of the masking and padding altogether), make sure you are on version `2.5` and import as follows
+
+```python
+import torch
+from vit_pytorch.na_vit_nested_tensor import NaViT
+
+v = NaViT(
+    image_size = 256,
+    patch_size = 32,
+    num_classes = 1000,
+    dim = 1024,
+    depth = 6,
+    heads = 16,
+    mlp_dim = 2048,
+    dropout = 0.,
+    emb_dropout = 0.,
+    token_dropout_prob = 0.1
+)
+
+# 5 images of different resolutions - List[Tensor]
+
+images = [
+    torch.randn(3, 256, 256), torch.randn(3, 128, 128),
+    torch.randn(3, 128, 256), torch.randn(3, 256, 128),
+    torch.randn(3, 64, 256)
+]
+
+preds = v(images)
+
+assert preds.shape == (5, 1000)
+```
+
 ## Distillation
 
 <img src="./images/distill.png" width="300px"></img>
@@ -754,6 +805,38 @@ img = torch.randn(1, 3, 256, 256)
 pred = mbvit_xs(img) # (1, 1000)
 ```
 
+## XCiT
+
+<img src="./images/xcit.png" width="400px"></img>
+
+This <a href="https://arxiv.org/abs/2106.09681">paper</a> introduces the cross covariance attention (abbreviated XCA). One can think of it as doing attention across the features dimension rather than the spatial one (another perspective would be a dynamic 1x1 convolution, the kernel being attention map defined by spatial correlations).
+
+Technically, this amounts to simply transposing the query, key, values before executing cosine similarity attention with learned temperature.
+
+```python
+import torch
+from vit_pytorch.xcit import XCiT
+
+v = XCiT(
+    image_size = 256,
+    patch_size = 32,
+    num_classes = 1000,
+    dim = 1024,
+    depth = 12,                     # depth of xcit transformer
+    cls_depth = 2,                  # depth of cross attention of CLS tokens to patch, attention pool at end
+    heads = 16,
+    mlp_dim = 2048,
+    dropout = 0.1,
+    emb_dropout = 0.1,
+    layer_dropout = 0.05,           # randomly dropout 5% of the layers
+    local_patch_kernel_size = 3     # kernel size of the local patch interaction module (depthwise convs)
+)
+
+img = torch.randn(1, 3, 256, 256)
+
+preds = v(img) # (1, 1000)
+```
+
 ## Simple Masked Image Modeling
 
 <img src="./images/simmim.png" width="400px"/>
@@ -889,7 +972,7 @@ torch.save(model.state_dict(), './pretrained-net.pt')
 
 <img src="./images/mp3.png" width="400px"></img>
 
-New <a href="https://arxiv.org/abs/2207.07611">paper</a> that introduces masked position prediction pre-training criteria. This strategy is more efficient than the Masked Autoencoder strategy and has comparable performance.  
+New <a href="https://arxiv.org/abs/2207.07611">paper</a> that introduces masked position prediction pre-training criteria. This strategy is more efficient than the Masked Autoencoder strategy and has comparable performance.
 
 ```python
 import torch
@@ -1135,7 +1218,8 @@ pred = cct(video)
 
 <img src="./images/vivit.png" width="350px"></img>
 
-This <a href="https://arxiv.org/abs/2103.15691">paper</a> offers 3 different types of architectures for efficient attention of videos, with the main theme being factorizing the attention across space and time. This repository will offer the first variant, which is a spatial transformer followed by a temporal one.
+This <a href="https://arxiv.org/abs/2103.15691">paper</a> offers 3 different types of architectures for efficient attention of videos, with the main theme being factorizing the attention across space and time. This repository includes the factorized encoder and the factorized self-attention variant.
+The factorized encoder variant is a spatial transformer followed by a temporal one. The factorized self-attention variant is a spatio-temporal transformer with alternating spatial and temporal self-attention layers.
 
 ```python
 import torch
@@ -1151,7 +1235,8 @@ v = ViT(
     spatial_depth = 6,         # depth of the spatial transformer
     temporal_depth = 6,        # depth of the temporal transformer
     heads = 8,
-    mlp_dim = 2048
+    mlp_dim = 2048,
+    variant = 'factorized_encoder', # or 'factorized_self_attention'
 )
 
 video = torch.randn(4, 3, 16, 128, 128) # (batch, channels, frames, height, width)
@@ -1273,10 +1358,10 @@ learner = Dino(
     hidden_layer = 'to_latent',        # hidden layer name or index, from which to extract the embedding
     projection_hidden_size = 256,      # projector network hidden dimension
     projection_layers = 4,             # number of layers in projection network
-    num_classes_K = 65336,             # output logits dimensions (referenced as K in paper)
+    num_classes_K = 65536,             # output logits dimensions (referenced as K in paper)
     student_temp = 0.9,                # student temperature
     teacher_temp = 0.04,               # teacher temperature, needs to be annealed from 0.04 to 0.07 over 30 epochs
-    local_upper_crop_scale = 0.4,      # upper bound for local crop - 0.4 was recommended in the paper 
+    local_upper_crop_scale = 0.4,      # upper bound for local crop - 0.4 was recommended in the paper
     global_lower_crop_scale = 0.5,     # lower bound for global crop - 0.5 was recommended in the paper
     moving_average_decay = 0.9,        # moving average of encoder - paper showed anywhere from 0.9 to 0.999 was ok
     center_moving_average_decay = 0.9, # moving average of teacher centers - paper showed anywhere from 0.9 to 0.999 was ok
@@ -1490,6 +1575,105 @@ logits, embeddings = v(img)
 embeddings # ((1, 257, 192), (1, 17, 384)) - (batch x patches x dimension) <- large and small scales respectively
 ```
 
+## JetViT
+
+<img src="./images/jet_vit.png" width="400px"></img>
+
+Efficient high-resolution ViT that converts a pretrained full-attention ViT into a hybrid of ReLU-based linear attention with squeeze dynamic convolution, window attention, and a small number of full softmax attention layers via a three-stage Post-Training Attention Search.
+
+Credit goes to [Pranoy](https://github.com/pranoyr) for the implementation!
+
+Stage 1 — Each layer randomly picks between window and linear attention every forward pass.
+
+```python
+import torch
+from vit_pytorch.jet_vit import JetViT
+
+# LA - Linear Attention
+# WA - Window Attention
+# FA - Full Attention
+
+model = JetViT(
+    image_size = 224,
+    patch_size = 16,
+    num_classes = 1000,
+    dim = 512,
+    depth = 6,
+    heads = 8,
+    dim_head = 64,
+    mlp_dim = 3072,
+    window_size = 7,
+    attn_layers = [('WA', 'LA')] * 6,   # randomly sample between window and linear attention
+)
+```
+
+Stage 2 — Each layer randomly picks between its stage-1 winner and full softmax attention.
+
+```python
+model = JetViT(
+    image_size = 224,
+    patch_size = 16,
+    num_classes = 1000,
+    dim = 512,
+    depth = 6,
+    heads = 8,
+    dim_head = 64,
+    mlp_dim = 3072,
+    window_size = 7,
+    attn_layers = [('LA', 'FA'), ('LA', 'FA'), ('WA', 'FA'), ('LA', 'FA'), ('LA', 'FA'), ('WA', 'FA')]
+)
+```
+
+Inference: every layer uses its winning attention type deterministically.
+
+```python
+model = JetViT(
+    image_size = 224,
+    patch_size = 16,
+    num_classes = 1000,
+    dim = 512,
+    depth = 6,
+    heads = 8,
+    dim_head = 64,
+    mlp_dim = 3072,
+    window_size = 7,
+    attn_layers = ['LA', 'LA', 'WA', 'FA', 'LA', 'FA']
+)
+
+img = torch.randn(1, 3, 224, 224)
+out = model(img)  # (1, 1000)
+```
+
+## ViT-5
+
+<img src="./images/vit-5.png" width="400px"></img>
+
+Vision Transformers for the Mid-2020s. An updated ViT architecture combining RMSNorm, QK-norm, LayerScale, register tokens, and 2D axial RoPE (applied to patch and register tokens, excluding the CLS token).
+
+Contributed by [Pranoy](https://github.com/pranoyr)!
+
+```python
+import torch
+from vit_pytorch.vit_5 import ViT
+
+model = ViT(
+    image_size    = 256,
+    patch_size    = 16,
+    num_classes   = 1000,
+    dim           = 512,
+    depth         = 6,
+    heads         = 8,
+    mlp_dim       = 1024,
+    dim_head      = 64,
+    dropout       = 0.1,
+    emb_dropout   = 0.1,
+    num_registers = 4,
+)
+
+img = torch.randn(2, 3, 256, 256)
+logits = model(img)  # (2, 1000)
+```
+
 ## Research Ideas
 
 ### Efficient Attention
@@ -1650,7 +1834,7 @@ Coming from computer vision and new to transformers? Here are some resources tha
 
 ```bibtex
 @misc{touvron2020training,
-    title   = {Training data-efficient image transformers & distillation through attention}, 
+    title   = {Training data-efficient image transformers & distillation through attention},
     author  = {Hugo Touvron and Matthieu Cord and Matthijs Douze and Francisco Massa and Alexandre Sablayrolles and Hervé Jégou},
     year    = {2020},
     eprint  = {2012.12877},
@@ -1683,7 +1867,7 @@ Coming from computer vision and new to transformers? Here are some resources tha
 
 ```bibtex
 @misc{touvron2021going,
-    title   = {Going deeper with Image Transformers}, 
+    title   = {Going deeper with Image Transformers},
     author  = {Hugo Touvron and Matthieu Cord and Alexandre Sablayrolles and Gabriel Synnaeve and Hervé Jégou},
     year    = {2021},
     eprint  = {2103.17239},
@@ -1716,7 +1900,7 @@ Coming from computer vision and new to transformers? Here are some resources tha
 
 ```bibtex
 @misc{heo2021rethinking,
-    title   = {Rethinking Spatial Dimensions of Vision Transformers}, 
+    title   = {Rethinking Spatial Dimensions of Vision Transformers},
     author  = {Byeongho Heo and Sangdoo Yun and Dongyoon Han and Sanghyuk Chun and Junsuk Choe and Seong Joon Oh},
     year    = {2021},
     eprint  = {2103.16302},
@@ -1760,7 +1944,7 @@ Coming from computer vision and new to transformers? Here are some resources tha
 
 ```bibtex
 @misc{su2021roformer,
-    title   = {RoFormer: Enhanced Transformer with Rotary Position Embedding}, 
+    title   = {RoFormer: Enhanced Transformer with Rotary Position Embedding},
     author  = {Jianlin Su and Yu Lu and Shengfeng Pan and Bo Wen and Yunfeng Liu},
     year    = {2021},
     eprint  = {2104.09864},
@@ -1782,7 +1966,7 @@ Coming from computer vision and new to transformers? Here are some resources tha
 
 ```bibtex
 @misc{chen2021regionvit,
-    title   = {RegionViT: Regional-to-Local Attention for Vision Transformers}, 
+    title   = {RegionViT: Regional-to-Local Attention for Vision Transformers},
     author  = {Chun-Fu Chen and Rameswar Panda and Quanfu Fan},
     year    = {2021},
     eprint  = {2106.02689},
@@ -1793,7 +1977,7 @@ Coming from computer vision and new to transformers? Here are some resources tha
 
 ```bibtex
 @misc{wang2021crossformer,
-    title   = {CrossFormer: A Versatile Vision Transformer Hinging on Cross-scale Attention}, 
+    title   = {CrossFormer: A Versatile Vision Transformer Hinging on Cross-scale Attention},
     author  = {Wenxiao Wang and Lu Yao and Long Chen and Binbin Lin and Deng Cai and Xiaofei He and Wei Liu},
     year    = {2021},
     eprint  = {2108.00154},
@@ -1815,7 +1999,7 @@ Coming from computer vision and new to transformers? Here are some resources tha
 
 ```bibtex
 @misc{he2021masked,
-    title   = {Masked Autoencoders Are Scalable Vision Learners}, 
+    title   = {Masked Autoencoders Are Scalable Vision Learners},
     author  = {Kaiming He and Xinlei Chen and Saining Xie and Yanghao Li and Piotr Dollár and Ross Girshick},
     year    = {2021},
     eprint  = {2111.06377},
@@ -1826,7 +2010,7 @@ Coming from computer vision and new to transformers? Here are some resources tha
 
 ```bibtex
 @misc{xie2021simmim,
-    title   = {SimMIM: A Simple Framework for Masked Image Modeling}, 
+    title   = {SimMIM: A Simple Framework for Masked Image Modeling},
     author  = {Zhenda Xie and Zheng Zhang and Yue Cao and Yutong Lin and Jianmin Bao and Zhuliang Yao and Qi Dai and Han Hu},
     year    = {2021},
     eprint  = {2111.09886},
@@ -1859,7 +2043,7 @@ Coming from computer vision and new to transformers? Here are some resources tha
 
 ```bibtex
 @misc{lee2021vision,
-    title   = {Vision Transformer for Small-Size Datasets}, 
+    title   = {Vision Transformer for Small-Size Datasets},
     author  = {Seung Hoon Lee and Seunghyun Lee and Byung Cheol Song},
     year    = {2021},
     eprint  = {2112.13492},
@@ -1881,7 +2065,7 @@ Coming from computer vision and new to transformers? Here are some resources tha
 
 ```bibtex
 @misc{yang2022scalablevit,
-    title   = {ScalableViT: Rethinking the Context-oriented Generalization of Vision Transformer}, 
+    title   = {ScalableViT: Rethinking the Context-oriented Generalization of Vision Transformer},
     author  = {Rui Yang and Hailong Ma and Jie Wu and Yansong Tang and Xuefeng Xiao and Min Zheng and Xiu Li},
     year    = {2022},
     eprint  = {2203.10790},
@@ -2002,4 +2186,293 @@ Coming from computer vision and new to transformers? Here are some resources tha
 }
 ```
 
+```bibtex
+@inproceedings{Darcet2023VisionTN,
+    title   = {Vision Transformers Need Registers},
+    author  = {Timoth'ee Darcet and Maxime Oquab and Julien Mairal and Piotr Bojanowski},
+    year    = {2023},
+    url     = {https://api.semanticscholar.org/CorpusID:263134283}
+}
+```
+
+```bibtex
+@inproceedings{ElNouby2021XCiTCI,
+    title   = {XCiT: Cross-Covariance Image Transformers},
+    author  = {Alaaeldin El-Nouby and Hugo Touvron and Mathilde Caron and Piotr Bojanowski and Matthijs Douze and Armand Joulin and Ivan Laptev and Natalia Neverova and Gabriel Synnaeve and Jakob Verbeek and Herv{\'e} J{\'e}gou},
+    booktitle = {Neural Information Processing Systems},
+    year    = {2021},
+    url     = {https://api.semanticscholar.org/CorpusID:235458262}
+}
+```
+
+```bibtex
+@inproceedings{Koner2024LookupViTCV,
+    title   = {LookupViT: Compressing visual information to a limited number of tokens},
+    author  = {Rajat Koner and Gagan Jain and Prateek Jain and Volker Tresp and Sujoy Paul},
+    year    = {2024},
+    url     = {https://api.semanticscholar.org/CorpusID:271244592}
+}
+```
+
+```bibtex
+@article{Bao2022AllAW,
+    title   = {All are Worth Words: A ViT Backbone for Diffusion Models},
+    author  = {Fan Bao and Shen Nie and Kaiwen Xue and Yue Cao and Chongxuan Li and Hang Su and Jun Zhu},
+    journal = {2023 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
+    year    = {2022},
+    pages   = {22669-22679},
+    url     = {https://api.semanticscholar.org/CorpusID:253581703}
+}
+```
+
+```bibtex
+@misc{Rubin2024,
+    author  = {Ohad Rubin},
+    url     = {https://medium.com/@ohadrubin/exploring-weight-decay-in-layer-normalization-challenges-and-a-reparameterization-solution-ad4d12c24950}
+}
+```
+
+```bibtex
+@inproceedings{Loshchilov2024nGPTNT,
+    title   = {nGPT: Normalized Transformer with Representation Learning on the Hypersphere},
+    author  = {Ilya Loshchilov and Cheng-Ping Hsieh and Simeng Sun and Boris Ginsburg},
+    year    = {2024},
+    url     = {https://api.semanticscholar.org/CorpusID:273026160}
+}
+```
+
+```bibtex
+@inproceedings{Liu2017DeepHL,
+    title   = {Deep Hyperspherical Learning},
+    author  = {Weiyang Liu and Yanming Zhang and Xingguo Li and Zhen Liu and Bo Dai and Tuo Zhao and Le Song},
+    booktitle = {Neural Information Processing Systems},
+    year    = {2017},
+    url     = {https://api.semanticscholar.org/CorpusID:5104558}
+}
+```
+
+```bibtex
+@inproceedings{Zhou2024ValueRL,
+    title   = {Value Residual Learning For Alleviating Attention Concentration In Transformers},
+    author  = {Zhanchao Zhou and Tianyi Wu and Zhiyun Jiang and Zhenzhong Lan},
+    year    = {2024},
+    url     = {https://api.semanticscholar.org/CorpusID:273532030}
+}
+```
+
+```bibtex
+@article{Zhu2024HyperConnections,
+    title   = {Hyper-Connections},
+    author  = {Defa Zhu and Hongzhi Huang and Zihao Huang and Yutao Zeng and Yunyao Mao and Banggu Wu and Qiyang Min and Xun Zhou},
+    journal = {ArXiv},
+    year    = {2024},
+    volume  = {abs/2409.19606},
+    url     = {https://api.semanticscholar.org/CorpusID:272987528}
+}
+```
+
+```bibtex
+@inproceedings{Fuller2025SimplerFV,
+    title   = {Simpler Fast Vision Transformers with a Jumbo CLS Token},
+    author  = {Anthony Fuller and Yousef Yassin and Daniel G. Kyrollos and Evan Shelhamer and James R. Green},
+    year    = {2025},
+    url     = {https://api.semanticscholar.org/CorpusID:276557720}
+}
+```
+
+```bibtex
+@misc{xiong2025ndrope,
+    author = {Jerry Xiong},
+    title  = {On n-dimensional rotary positional embeddings},
+    year   = {2025},
+    url    = {https://jerryxio.ng/posts/nd-rope/}
+}
+```
+
+```bibtex
+@inproceedings{anonymous2025vat,
+    title   = {{VAT}: Vision Action Transformer by Unlocking Full Representation of ViT},
+    author  = {Anonymous},
+    booktitle = {Submitted to The Fourteenth International Conference on Learning Representations},
+    year    = {2025},
+    url     = {https://openreview.net/forum?id=TalHOvvLZu},
+    note    = {under review}
+}
+```
+
+```bibtex
+@misc{carrigg2025decorrelationspeedsvisiontransformers,
+    title   = {Decorrelation Speeds Up Vision Transformers},
+    author  = {Kieran Carrigg and Rob van Gastel and Melda Yeghaian and Sander Dalm and Faysal Boughorbel and Marcel van Gerven},
+    year    = {2025},
+    eprint  = {2510.14657},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.CV},
+    url     = {https://arxiv.org/abs/2510.14657},
+}
+```
+
+```bibtex
+@misc{gopalakrishnan2025decouplingwhatwherepolar,
+    title   = {Decoupling the "What" and "Where" With Polar Coordinate Positional Embeddings},
+    author  = {Anand Gopalakrishnan and Robert Csordás and Jürgen Schmidhuber and Michael C. Mozer},
+    year    = {2025},
+    eprint  = {2509.10534},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.LG},
+    url     = {https://arxiv.org/abs/2509.10534},
+}
+```
+
+```bibtex
+@misc{qiu2025gatedattentionlargelanguage,
+    title   = {Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free},
+    author  = {Zihan Qiu and Zekun Wang and Bo Zheng and Zeyu Huang and Kaiyue Wen and Songlin Yang and Rui Men and Le Yu and Fei Huang and Suozhi Huang and Dayiheng Liu and Jingren Zhou and Junyang Lin},
+    year    = {2025},
+    eprint  = {2505.06708},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.CL},
+    url     = {https://arxiv.org/abs/2505.06708}
+}
+```
+
+```bibtex
+@misc{chen2026postlayernormbackstableexpressive,
+    title   = {Post-LayerNorm Is Back: Stable, ExpressivE, and Deep},
+    author  = {Chen Chen and Lai Wei},
+    year    = {2026},
+    eprint  = {2601.19895},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.LG},
+    url     = {https://arxiv.org/abs/2601.19895},
+}
+```
+
+```bibtex
+@misc{intelligence2025pi06vlalearnsexperience,
+      title     = {$\pi^{*}_{0.6}$: a VLA That Learns From Experience},
+      author    = {Physical Intelligence and Ali Amin and Raichelle Aniceto and Ashwin Balakrishna and Kevin Black and Ken Conley and Grace Connors and James Darpinian and Karan Dhabalia and Jared DiCarlo and Danny Driess and Michael Equi and Adnan Esmail and Yunhao Fang and Chelsea Finn and Catherine Glossop and Thomas Godden and Ivan Goryachev and Lachy Groom and Hunter Hancock and Karol Hausman and Gashon Hussein and Brian Ichter and Szymon Jakubczak and Rowan Jen and Tim Jones and Ben Katz and Liyiming Ke and Chandra Kuchi and Marinda Lamb and Devin LeBlanc and Sergey Levine and Adrian Li-Bell and Yao Lu and Vishnu Mano and Mohith Mothukuri and Suraj Nair and Karl Pertsch and Allen Z. Ren and Charvi Sharma and Lucy Xiaoyang Shi and Laura Smith and Jost Tobias Springenberg and Kyle Stachowicz and Will Stoeckle and Alex Swerdlow and James Tanner and Marcel Torne and Quan Vuong and Anna Walling and Haohuan Wang and Blake Williams and Sukwon Yoo and Lili Yu and Ury Zhilinsky and Zhiyuan Zhou},
+      year      = {2025},
+      eprint    = {2511.14759},
+      archivePrefix = {arXiv},
+      primaryClass = {cs.LG},
+      url       = {https://arxiv.org/abs/2511.14759},
+}
+```
+
+```bibtex
+@misc{kimiteam2026attentionresiduals,
+    title   = {Attention Residuals},
+    author  = {Kimi Team and Guangyu Chen and Yu Zhang and Jianlin Su and Weixin Xu and Siyuan Pan and Yaoyu Wang and Yucheng Wang and Guanduo Chen and Bohong Yin and Yutian Chen and Junjie Yan and Ming Wei and Y. Zhang and Fanqing Meng and Chao Hong and Xiaotong Xie and Shaowei Liu and Enzhe Lu and Yunpeng Tai and Yanru Chen and Xin Men and Haiqing Guo and Y. Charles and Haoyu Lu and Lin Sui and Jinguo Zhu and Zaida Zhou and Weiran He and Weixiao Huang and Xinran Xu and Yuzhi Wang and Guokun Lai and Yulun Du and Yuxin Wu and Zhilin Yang and Xinyu Zhou},
+    year    = {2026},
+    eprint  = {2603.15031},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.CL},
+    url     = {https://arxiv.org/abs/2603.15031},
+}
+```
+
+```bibtex
+@misc{balestriero2025lejepa,
+    title   = {LeJEPA: Provable and Scalable Self-Supervised Learning Without the Heuristics},
+    author  = {Randall Balestriero and Yann LeCun},
+    year    = {2025},
+    eprint  = {2511.08544},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.LG},
+    url     = {https://arxiv.org/abs/2511.08544},
+}
+```
+
+```bibtex
+@misc{oh2026revisitingresidualconnectionsorthogonal,
+    title   = {Revisiting Residual Connections: Orthogonal Updates for Stable and Efficient Deep Networks},
+    author  = {Giyeong Oh and Woohyun Cho and Siyeol Kim and Suhwan Choi and Youngjae Yu},
+    year    = {2026},
+    eprint  = {2505.11881},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.CV},
+    url     = {https://arxiv.org/abs/2505.11881},
+}
+```
+
+```bibtex
+@inproceedings{niu2026learning,
+    title   = {Learning to Grasp Anything By Playing with Random Toys},
+    author  = {Dantong Niu and Yuvan Sharma and Baifeng Shi and Rachel Ding and Matteo Gioia and Haoru Xue and Henry Tsai and Konstantinos Kallidromitis and Anirudh Pai and S. Shankar Sastry and Trevor Darrell and Jitendra Malik and Roei Herzig},
+    booktitle = {The Fourteenth International Conference on Learning Representations},
+    year    = {2026},
+    url     = {https://openreview.net/forum?id=NZDaMcpXZm}
+}
+```
+
+```bibtex
+@misc{marouani2026revisitingclspatchtoken,
+    title   = {Revisiting [CLS] and Patch Token Interaction in Vision Transformers},
+    author  = {Alexis Marouani and Oriane Siméoni and Hervé Jégou and Piotr Bojanowski and Huy V. Vo},
+    year    = {2026},
+    eprint  = {2602.08626},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.CV},
+    url     = {https://arxiv.org/abs/2602.08626},
+}
+```
+
+```bibtex
+@misc{kim2026exploring,
+    title   = {Exploring High-Order Self-Similarity for Video Understanding},
+    author  = {Manjin Kim and Heeseung Kwon and Karteek Alahari and Minsu Cho},
+    year    = {2026},
+    url     = {https://openreview.net/forum?id=Co6SCyBIjo}
+}
+```
+
+```bibtex
+@misc{akbar2026weaksigregcovarianceregularizationstable,
+    title   = {Weak-SIGReg: Covariance Regularization for Stable Deep Learning},
+    author  = {Habibullah Akbar},
+    year    = {2026},
+    eprint  = {2603.05924},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.LG},
+    url     = {https://arxiv.org/abs/2603.05924},
+}
+```
+
+```bibtex
+@misc{zou2026jetvitefficienthighresolutionvision,
+    title   = {JetViT: Efficient High-Resolution Vision Transformer with Post-Training Attention Search},
+    author  = {Dongyun Zou and Zhuoyang Zhang and Junyu Chen and Wenkun He and Qinhe Peng and Hanrong Ye and Yao Lu and Hongxu Yin and Yu Wang and Song Han and Han Cai},
+    year    = {2026},
+    eprint  = {2605.26636},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.CV},
+    url     = {https://arxiv.org/abs/2605.26636},
+}
+```
+
+```bibtex
+@misc{yoshihashi2026whatwheretransformer,
+    title   = {What-Where Transformer: A Slot-Centric Visual Backbone for Concurrent Representation and Localization},
+    author  = {Ryota Yoshihashi and Masahiro Kada and Satoshi Ikehata and Rei Kawakami and Ikuro Sato},
+    year    = {2026},
+    eprint  = {2605.12021},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.CV},
+    url     = {https://arxiv.org/abs/2605.12021},
+}
+```
+
+```bibtex
+@misc{wang2026vit5,
+    title   = {ViT-5: Vision Transformers for The Mid-2020s},
+    author  = {Feng Wang and Sucheng Ren and Tiezheng Zhang and Predrag Neskovic and Anand Bhattad and Cihang Xie and Alan Yuille},
+    year    = {2026},
+    eprint  = {2602.08071},
+    archivePrefix = {arXiv},
+    primaryClass  = {cs.CV},
+    url     = {https://arxiv.org/abs/2602.08071},
+}
+```
+
 *I visualise a time when we will be to robots what dogs are to humans, and I’m rooting for the machines.* — Claude Shannon

pyproject.toml

@@ -0,0 +1,63 @@
+[build-system]
+requires = ["setuptools>=61", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "vit-pytorch"
+version = "1.24.1"
+description = "Vision Transformer (ViT) - Pytorch"
+readme = { file = "README.md", content-type = "text/markdown" }
+license = { file = "LICENSE" }
+authors = [
+  { name = "Phil Wang", email = "lucidrains@gmail.com" },
+]
+requires-python = ">=3.8"
+keywords = [
+  "artificial intelligence",
+  "attention mechanism",
+  "image recognition",
+]
+classifiers = [
+  "Development Status :: 4 - Beta",
+  "Intended Audience :: Developers",
+  "Topic :: Scientific/Engineering :: Artificial Intelligence",
+  "License :: OSI Approved :: MIT License",
+  "Programming Language :: Python :: 3",
+  "Programming Language :: Python :: 3 :: Only",
+  "Programming Language :: Python :: 3.8",
+  "Programming Language :: Python :: 3.9",
+  "Programming Language :: Python :: 3.10",
+  "Programming Language :: Python :: 3.11",
+  "Programming Language :: Python :: 3.12",
+]
+dependencies = [
+  "einops>=0.8.2",
+  "torch>=2.4",
+  "torchvision",
+]
+
+[project.optional-dependencies]
+test = [
+  "pytest",
+  "torch==2.4.0",
+  "torchvision==0.19.0",
+]
+
+[project.urls]
+Homepage = "https://codeberg.org/lucidrains/vit-pytorch"
+Repository = "https://codeberg.org/lucidrains/vit-pytorch"
+
+[tool.setuptools]
+include-package-data = true
+
+[tool.setuptools.packages.find]
+include = ["vit_pytorch*"]
+exclude = ["examples*", "tests*", "test*"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests", "."]
+python_files = ["test_*.py", "*_test.py"]
+addopts = "-q"
+filterwarnings = [
+  "ignore::FutureWarning",
+]

setup.py

@@ -1,38 +0,0 @@
-from setuptools import setup, find_packages
-
-setup(
-  name = 'vit-pytorch',
-  packages = find_packages(exclude=['examples']),
-  version = '1.2.8',
-  license='MIT',
-  description = 'Vision Transformer (ViT) - Pytorch',
-  long_description_content_type = 'text/markdown',
-  author = 'Phil Wang',
-  author_email = 'lucidrains@gmail.com',
-  url = 'https://github.com/lucidrains/vit-pytorch',
-  keywords = [
-    'artificial intelligence',
-    'attention mechanism',
-    'image recognition'
-  ],
-  install_requires=[
-    'einops>=0.6.1',
-    'torch>=1.10',
-    'torchvision'
-  ],
-  setup_requires=[
-    'pytest-runner',
-  ],
-  tests_require=[
-    'pytest',
-    'torch==1.12.1',
-    'torchvision==0.13.1'
-  ],
-  classifiers=[
-    'Development Status :: 4 - Beta',
-    'Intended Audience :: Developers',
-    'Topic :: Scientific/Engineering :: Artificial Intelligence',
-    'License :: OSI Approved :: MIT License',
-    'Programming Language :: Python :: 3.6',
-  ],
-)

tests/test_vit.py

@@ -1,7 +1,7 @@
 import torch
 from vit_pytorch import ViT
 
-def test():
+def test_vit():
     v = ViT(
         image_size = 256,
         patch_size = 32,

train_vit_decorr.py

@@ -0,0 +1,111 @@
+# /// script
+# dependencies = [
+#   "accelerate",
+#   "vit-pytorch",
+#   "wandb"
+# ]
+# ///
+
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+
+import torchvision.transforms as T
+from torchvision.datasets import CIFAR100
+
+# constants
+
+BATCH_SIZE = 32
+LEARNING_RATE = 3e-4
+EPOCHS = 10
+DECORR_LOSS_WEIGHT = 1e-1
+
+TRACK_EXPERIMENT_ONLINE = False
+
+# helpers
+
+def exists(v):
+    return v is not None
+
+# data
+
+transform = T.Compose([
+    T.ToTensor(),
+    T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
+])
+
+dataset = CIFAR100(
+    root = 'data',
+    download = True,
+    train = True,
+    transform = transform
+)
+
+dataloader = DataLoader(dataset, batch_size = BATCH_SIZE, shuffle = True)
+
+# model
+
+from vit_pytorch.vit_with_decorr import ViT
+
+vit = ViT(
+    dim = 128,
+    num_classes = 100,
+    image_size = 32,
+    patch_size = 4,
+    depth = 6,
+    heads = 8,
+    dim_head = 64,
+    mlp_dim = 128 * 4,
+    decorr_sample_frac = 1.,        # use all tokens
+    decorr_use_subspace = False,    # use subspace
+    decorr_dim_subspace = 64,
+    decorr_num_subspaces = 4,
+    decorr_layer_outputs_across_depth = False
+)
+
+# optim
+
+from torch.optim import Adam
+
+optim = Adam(vit.parameters(), lr = LEARNING_RATE)
+
+# prepare
+
+from accelerate import Accelerator
+
+accelerator = Accelerator()
+
+vit, optim, dataloader = accelerator.prepare(vit, optim, dataloader)
+
+# experiment
+
+import wandb
+
+wandb.init(
+    project = 'vit-decorr',
+    mode = 'disabled' if not TRACK_EXPERIMENT_ONLINE else 'online'
+)
+
+wandb.run.name = 'baseline'
+
+# loop
+
+for _ in range(EPOCHS):
+    for images, labels in dataloader:
+
+        logits, decorr_aux_loss = vit(images)
+        loss = F.cross_entropy(logits, labels)
+
+
+        total_loss = (
+            loss +
+            decorr_aux_loss * DECORR_LOSS_WEIGHT
+        )
+
+        wandb.log(dict(loss = loss, decorr_loss = decorr_aux_loss))
+
+        accelerator.print(f'loss: {loss.item():.3f} | decorr aux loss: {decorr_aux_loss.item():.3f}')
+
+        accelerator.backward(total_loss)
+        optim.step()
+        optim.zero_grad()

vit_pytorch/__init__.py

@@ -1,10 +1,3 @@
-import torch
-from packaging import version
-
-if version.parse(torch.__version__) >= version.parse('2.0.0'):
-    from einops._torch_specific import allow_ops_in_compiled_graph
-    allow_ops_in_compiled_graph()
-
 from vit_pytorch.vit import ViT
 from vit_pytorch.simple_vit import SimpleViT
 

vit_pytorch/accept_video_wrapper.py

@@ -0,0 +1,230 @@
+from __future__ import annotations
+
+from contextlib import nullcontext
+
+import torch
+from torch import is_tensor, randn
+from torch.nn import Module, Linear, Parameter
+from torch.utils._pytree import tree_flatten, tree_unflatten
+
+from vit_pytorch.vivit_with_moss import MOSS
+
+from einops import rearrange, repeat
+
+# helper functions
+
+def exists(v):
+    return v is not None
+
+def default(v, d):
+    return v if exists(v) else d
+
+def pair(t):
+    return t if isinstance(t, tuple) else (t, t)
+
+# classes
+
+class AcceptVideoWrapper(Module):
+    def __init__(
+        self,
+        image_net: Module,
+        forward_function = 'forward',
+        add_time_pos_emb = False,
+        dim_emb = None,
+        time_seq_len = None,
+        embed_is_channel_first = False,
+        output_pos_add_pos_emb = 0, # defaults to first output position to add embedding
+        proj_embed_to_dim = None,
+        patch_size = None,
+        moss: Module | dict | None = None
+    ):
+        super().__init__()
+        self.image_net = image_net
+        self.forward_function = forward_function # for openclip, used in TRI-LBM
+
+        self.add_time_pos_emb = add_time_pos_emb
+        self.output_pos_add_pos_emb = output_pos_add_pos_emb
+
+        # maybe project the image embedding
+
+        self.embed_proj = None
+
+        if exists(proj_embed_to_dim):
+            assert exists(dim_emb), '`dim_emb` must be passed in'
+            self.embed_proj = Linear(dim_emb, proj_embed_to_dim)
+
+        # time positional embedding
+
+        if add_time_pos_emb:
+            assert exists(dim_emb) and exists(time_seq_len), '`dim_emb` and `time_seq_len` must be set if adding positional embeddings to the output'
+            self.time_seq_len = time_seq_len
+
+            dim_pos_emb = default(proj_embed_to_dim, dim_emb)
+
+            self.pos_emb = Parameter(randn(time_seq_len, dim_pos_emb) * 1e-2)
+
+        self.embed_is_channel_first = embed_is_channel_first
+
+        # patch size and moss
+
+        if not exists(patch_size):
+            if hasattr(image_net, 'patch_size'):
+                patch_size = image_net.patch_size
+            elif hasattr(image_net, 'vit') and hasattr(image_net.vit, 'patch_size'):
+                patch_size = image_net.vit.patch_size
+
+        self.patch_size = patch_size
+
+        if isinstance(moss, dict):
+            moss = MOSS(**moss)
+
+        self.moss = moss
+
+        if exists(self.moss):
+            assert exists(self.patch_size), '`patch_size` must be provided either on the `image_net` or passed in explicitly if using MOSS'
+
+    def forward(
+        self,
+        video, # (b c t h w)
+        eval_with_no_grad = False,
+        forward_kwargs = dict()
+    ):
+        add_time_pos_emb = self.add_time_pos_emb
+        time = video.shape[2]
+
+        # maybe validate time positional embedding
+
+        if add_time_pos_emb:
+            assert time <= self.time_seq_len, f'received video with {time} frames but `time_seq_len` ({self.time_seq_len}) is too low'
+
+        video_height, video_width = video.shape[-2:]
+
+        video = rearrange(video, 'b c t h w -> b t c h w')
+
+        video = rearrange(video, 'b t ... -> (b t) ...')
+
+        # forward through image net for outputs
+
+        func = getattr(self.image_net, self.forward_function)
+
+        if eval_with_no_grad:
+            self.image_net.eval()
+
+        context = torch.no_grad if eval_with_no_grad else nullcontext
+
+        with context():
+            outputs = func(video, **forward_kwargs)
+
+        # handle multiple outputs, say logits and embeddings returned from extractor - also handle some reduce aux loss being returned
+
+        outputs, tree_spec = tree_flatten(outputs)
+
+        outputs = tuple(rearrange(t, '(b t) ... -> b t ...', t = time) if is_tensor(t) and t.numel() > 1 else t for t in outputs)
+
+        # maybe project embedding
+
+        if exists(self.embed_proj):
+            outputs = list(outputs)
+
+            embed = outputs[self.output_pos_add_pos_emb]
+
+            outputs[self.output_pos_add_pos_emb] = self.embed_proj(embed)
+
+        # maybe add time positional embedding
+
+        if add_time_pos_emb:
+
+            outputs = list(outputs)
+            embed = outputs[self.output_pos_add_pos_emb]
+
+            pos_emb = rearrange(self.pos_emb, 't d -> 1 t d')
+
+            # handle the network outputting embeddings with spatial dimensions intact - assume embedded dimension is last
+
+            dims_to_unsqueeze = embed.ndim - pos_emb.ndim
+
+            one_dims = ((1,) * dims_to_unsqueeze)
+
+            if self.embed_is_channel_first:
+                pos_emb = pos_emb.reshape(*pos_emb.shape, *one_dims)
+            else:
+                pos_emb = pos_emb.reshape(*pos_emb.shape[:2], *one_dims, pos_emb.shape[-1])
+
+            pos_emb = pos_emb[:, :embed.shape[1]]
+
+            embed = embed + pos_emb
+
+            outputs[self.output_pos_add_pos_emb] = embed
+
+        # moss - stack of ssts
+        # https://openreview.net/forum?id=Co6SCyBIjo
+
+        if exists(self.moss):
+            outputs = list(outputs)
+            embed = outputs[self.output_pos_add_pos_emb]
+
+            patch_h, patch_w = pair(self.patch_size)
+            num_h, num_w = video_height // patch_h, video_width // patch_w
+            num_patches = num_h * num_w
+
+            num_cls_tokens = embed.shape[-2] - num_patches
+            cls_tokens, patch_tokens = embed[:, :, :num_cls_tokens], embed[:, :, num_cls_tokens:]
+
+            patch_tokens = rearrange(patch_tokens, 'b t (h w) d -> b t h w d', h = num_h, w = num_w)
+            patch_tokens = self.moss(patch_tokens)
+            patch_tokens = rearrange(patch_tokens, 'b t h w d -> b t (h w) d')
+
+            embed = torch.cat((cls_tokens, patch_tokens), dim = -2)
+            outputs[self.output_pos_add_pos_emb] = embed
+
+        return tree_unflatten(outputs, tree_spec)
+
+# main
+
+if __name__ == '__main__':
+    from vit_pytorch import ViT
+
+    v = ViT(
+        image_size = 256,
+        patch_size = 32,
+        num_classes = 1000,
+        dim = 1024,
+        depth = 6,
+        heads = 16,
+        mlp_dim = 2048,
+        dropout = 0.1,
+        emb_dropout = 0.1
+    )
+
+    videos = torch.randn(1, 3, 7, 256, 256)
+
+    # step up the difficulty and return embeddings for robotics
+
+    from vit_pytorch.extractor import Extractor
+
+    v = Extractor(v)
+
+    moss_kwargs = dict(
+        dim = 512,
+        local_time = 3,
+        local_height = 3,
+        local_width = 3,
+        hidden_dim = 64,
+        orders = 2,
+        causal = True
+    )
+
+    video_acceptor = AcceptVideoWrapper(
+        v,
+        add_time_pos_emb = True,
+        output_pos_add_pos_emb = 1,
+        time_seq_len = 12,
+        dim_emb = 1024,
+        proj_embed_to_dim = 512,
+        moss = moss_kwargs
+    )
+
+    logits, embeddings = video_acceptor(videos, eval_with_no_grad = True) # always (batch, channels, time, height, width) - time is always dimension 2
+
+    assert logits.shape == (1, 7, 1000)
+    assert embeddings.shape == (1, 7, 65, 512)

vit_pytorch/ats_vit.py

@@ -110,18 +110,11 @@ class AdaptiveTokenSampling(nn.Module):
 
 # classes
 
-class PreNorm(nn.Module):
-    def __init__(self, dim, fn):
-        super().__init__()
-        self.norm = nn.LayerNorm(dim)
-        self.fn = fn
-    def forward(self, x, **kwargs):
-        return self.fn(self.norm(x), **kwargs)
-
 class FeedForward(nn.Module):
     def __init__(self, dim, hidden_dim, dropout = 0.):
         super().__init__()
         self.net = nn.Sequential(
+            nn.LayerNorm(dim),
             nn.Linear(dim, hidden_dim),
             nn.GELU(),
             nn.Dropout(dropout),
@@ -138,6 +131,7 @@ class Attention(nn.Module):
         self.heads = heads
         self.scale = dim_head ** -0.5
 
+        self.norm = nn.LayerNorm(dim)
         self.attend = nn.Softmax(dim = -1)
         self.dropout = nn.Dropout(dropout)
 
@@ -154,6 +148,7 @@ class Attention(nn.Module):
     def forward(self, x, *, mask):
         num_tokens = x.shape[1]
 
+        x = self.norm(x)
         qkv = self.to_qkv(x).chunk(3, dim = -1)
         q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
 
@@ -189,8 +184,8 @@ class Transformer(nn.Module):
         self.layers = nn.ModuleList([])
         for _, output_num_tokens in zip(range(depth), max_tokens_per_depth):
             self.layers.append(nn.ModuleList([
-                PreNorm(dim, Attention(dim, output_num_tokens = output_num_tokens, heads = heads, dim_head = dim_head, dropout = dropout)),
-                PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout))
+                Attention(dim, output_num_tokens = output_num_tokens, heads = heads, dim_head = dim_head, dropout = dropout),
+                FeedForward(dim, mlp_dim, dropout = dropout)
             ]))
 
     def forward(self, x):
@@ -221,7 +216,7 @@ class ViT(nn.Module):
     def __init__(self, *, image_size, patch_size, num_classes, dim, depth, max_tokens_per_depth, heads, mlp_dim, channels = 3, dim_head = 64, dropout = 0., emb_dropout = 0.):
         super().__init__()
         image_height, image_width = pair(image_size)
-        patch_height, patch_width = pair(patch_size)
+        self.patch_size = patch_height, patch_width = pair(patch_size)
 
         assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
 

vit_pytorch/cait.py

@@ -44,18 +44,11 @@ class LayerScale(nn.Module):
     def forward(self, x, **kwargs):
         return self.fn(x, **kwargs) * self.scale
 
-class PreNorm(nn.Module):
-    def __init__(self, dim, fn):
-        super().__init__()
-        self.norm = nn.LayerNorm(dim)
-        self.fn = fn
-    def forward(self, x, **kwargs):
-        return self.fn(self.norm(x), **kwargs)
-
 class FeedForward(nn.Module):
     def __init__(self, dim, hidden_dim, dropout = 0.):
         super().__init__()
         self.net = nn.Sequential(
+            nn.LayerNorm(dim),
             nn.Linear(dim, hidden_dim),
             nn.GELU(),
             nn.Dropout(dropout),
@@ -72,6 +65,7 @@ class Attention(nn.Module):
         self.heads = heads
         self.scale = dim_head ** -0.5
 
+        self.norm = nn.LayerNorm(dim)
         self.to_q = nn.Linear(dim, inner_dim, bias = False)
         self.to_kv = nn.Linear(dim, inner_dim * 2, bias = False)
 
@@ -89,6 +83,7 @@ class Attention(nn.Module):
     def forward(self, x, context = None):
         b, n, _, h = *x.shape, self.heads
 
+        x = self.norm(x)
         context = x if not exists(context) else torch.cat((x, context), dim = 1)
 
         qkv = (self.to_q(x), *self.to_kv(context).chunk(2, dim = -1))
@@ -115,8 +110,8 @@ class Transformer(nn.Module):
 
         for ind in range(depth):
             self.layers.append(nn.ModuleList([
-                LayerScale(dim, PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout)), depth = ind + 1),
-                LayerScale(dim, PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout)), depth = ind + 1)
+                LayerScale(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout), depth = ind + 1),
+                LayerScale(dim, FeedForward(dim, mlp_dim, dropout = dropout), depth = ind + 1)
             ]))
     def forward(self, x, context = None):
         layers = dropout_layers(self.layers, dropout = self.layer_dropout)

vit_pytorch/cct.py

@@ -316,6 +316,9 @@ class CCT(nn.Module):
         pooling_kernel_size=3,
         pooling_stride=2,
         pooling_padding=1,
+        dropout_rate=0.,
+        attention_dropout=0.1,
+        stochastic_depth_rate=0.1,
         *args, **kwargs
     ):
         super().__init__()
@@ -340,9 +343,9 @@ class CCT(nn.Module):
                                                            width=img_width),
             embedding_dim=embedding_dim,
             seq_pool=True,
-            dropout_rate=0.,
-            attention_dropout=0.1,
-            stochastic_depth=0.1,
+            dropout_rate=dropout_rate,
+            attention_dropout=attention_dropout,
+            stochastic_depth_rate=stochastic_depth_rate,
             *args, **kwargs)
 
     def forward(self, x):

vit_pytorch/cct_3d.py

@@ -167,8 +167,10 @@ class Tokenizer(nn.Module):
         stride,
         padding,
         frame_stride=1,
+        frame_padding=None,
         frame_pooling_stride=1,
         frame_pooling_kernel_size=1,
+        frame_pooling_padding=None,
         pooling_kernel_size=3,
         pooling_stride=2,
         pooling_padding=1,
@@ -188,16 +190,22 @@ class Tokenizer(nn.Module):
 
         n_filter_list_pairs = zip(n_filter_list[:-1], n_filter_list[1:])
 
+        if frame_padding is None:
+            frame_padding = frame_kernel_size // 2
+
+        if frame_pooling_padding is None:
+            frame_pooling_padding = frame_pooling_kernel_size // 2
+
         self.conv_layers = nn.Sequential(
             *[nn.Sequential(
                 nn.Conv3d(chan_in, chan_out,
                           kernel_size=(frame_kernel_size, kernel_size, kernel_size),
                           stride=(frame_stride, stride, stride),
-                          padding=(frame_kernel_size // 2, padding, padding), bias=conv_bias),
+                          padding=(frame_padding, padding, padding), bias=conv_bias),
                 nn.Identity() if not exists(activation) else activation(),
                 nn.MaxPool3d(kernel_size=(frame_pooling_kernel_size, pooling_kernel_size, pooling_kernel_size),
                              stride=(frame_pooling_stride, pooling_stride, pooling_stride),
-                             padding=(frame_pooling_kernel_size // 2, pooling_padding, pooling_padding)) if max_pool else nn.Identity()
+                             padding=(frame_pooling_padding, pooling_padding, pooling_padding)) if max_pool else nn.Identity()
             )
                 for chan_in, chan_out in n_filter_list_pairs
             ])
@@ -324,8 +332,10 @@ class CCT(nn.Module):
         n_conv_layers=1,
         frame_stride=1,
         frame_kernel_size=3,
+        frame_padding=None,
         frame_pooling_kernel_size=1,
         frame_pooling_stride=1,
+        frame_pooling_padding=None,
         kernel_size=7,
         stride=2,
         padding=3,
@@ -342,8 +352,10 @@ class CCT(nn.Module):
             n_output_channels=embedding_dim,
             frame_stride=frame_stride,
             frame_kernel_size=frame_kernel_size,
+            frame_padding=frame_padding,
             frame_pooling_stride=frame_pooling_stride,
             frame_pooling_kernel_size=frame_pooling_kernel_size,
+            frame_pooling_padding=frame_pooling_padding,
             kernel_size=kernel_size,
             stride=stride,
             padding=padding,

vit_pytorch/cross_vit.py

@@ -13,22 +13,13 @@ def exists(val):
 def default(val, d):
     return val if exists(val) else d
 
-# pre-layernorm
-
-class PreNorm(nn.Module):
-    def __init__(self, dim, fn):
-        super().__init__()
-        self.norm = nn.LayerNorm(dim)
-        self.fn = fn
-    def forward(self, x, **kwargs):
-        return self.fn(self.norm(x), **kwargs)
-
 # feedforward
 
 class FeedForward(nn.Module):
     def __init__(self, dim, hidden_dim, dropout = 0.):
         super().__init__()
         self.net = nn.Sequential(
+            nn.LayerNorm(dim),
             nn.Linear(dim, hidden_dim),
             nn.GELU(),
             nn.Dropout(dropout),
@@ -47,6 +38,7 @@ class Attention(nn.Module):
         self.heads = heads
         self.scale = dim_head ** -0.5
 
+        self.norm = nn.LayerNorm(dim)
         self.attend = nn.Softmax(dim = -1)
         self.dropout = nn.Dropout(dropout)
 
@@ -60,6 +52,7 @@ class Attention(nn.Module):
 
     def forward(self, x, context = None, kv_include_self = False):
         b, n, _, h = *x.shape, self.heads
+        x = self.norm(x)
         context = default(context, x)
 
         if kv_include_self:
@@ -86,8 +79,8 @@ class Transformer(nn.Module):
         self.norm = nn.LayerNorm(dim)
         for _ in range(depth):
             self.layers.append(nn.ModuleList([
-                PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout)),
-                PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout))
+                Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout),
+                FeedForward(dim, mlp_dim, dropout = dropout)
             ]))
 
     def forward(self, x):
@@ -121,8 +114,8 @@ class CrossTransformer(nn.Module):
         self.layers = nn.ModuleList([])
         for _ in range(depth):
             self.layers.append(nn.ModuleList([
-                ProjectInOut(sm_dim, lg_dim, PreNorm(lg_dim, Attention(lg_dim, heads = heads, dim_head = dim_head, dropout = dropout))),
-                ProjectInOut(lg_dim, sm_dim, PreNorm(sm_dim, Attention(sm_dim, heads = heads, dim_head = dim_head, dropout = dropout)))
+                ProjectInOut(sm_dim, lg_dim, Attention(lg_dim, heads = heads, dim_head = dim_head, dropout = dropout)),
+                ProjectInOut(lg_dim, sm_dim, Attention(sm_dim, heads = heads, dim_head = dim_head, dropout = dropout))
             ]))
 
     def forward(self, sm_tokens, lg_tokens):
@@ -177,12 +170,13 @@ class ImageEmbedder(nn.Module):
         dim,
         image_size,
         patch_size,
-        dropout = 0.
+        dropout = 0.,
+        channels = 3
     ):
         super().__init__()
         assert image_size % patch_size == 0, 'Image dimensions must be divisible by the patch size.'
         num_patches = (image_size // patch_size) ** 2
-        patch_dim = 3 * patch_size ** 2
+        patch_dim = channels * patch_size ** 2
 
         self.to_patch_embedding = nn.Sequential(
             Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_size, p2 = patch_size),
@@ -230,11 +224,12 @@ class CrossViT(nn.Module):
         cross_attn_dim_head = 64,
         depth = 3,
         dropout = 0.1,
-        emb_dropout = 0.1
+        emb_dropout = 0.1,
+        channels = 3
     ):
         super().__init__()
-        self.sm_image_embedder = ImageEmbedder(dim = sm_dim, image_size = image_size, patch_size = sm_patch_size, dropout = emb_dropout)
-        self.lg_image_embedder = ImageEmbedder(dim = lg_dim, image_size = image_size, patch_size = lg_patch_size, dropout = emb_dropout)
+        self.sm_image_embedder = ImageEmbedder(dim = sm_dim, channels= channels, image_size = image_size, patch_size = sm_patch_size, dropout = emb_dropout)
+        self.lg_image_embedder = ImageEmbedder(dim = lg_dim, channels = channels, image_size = image_size, patch_size = lg_patch_size, dropout = emb_dropout)
 
         self.multi_scale_encoder = MultiScaleEncoder(
             depth = depth,

vit_pytorch/cvt.py

@@ -34,19 +34,11 @@ class LayerNorm(nn.Module): # layernorm, but done in the channel dimension #1
         mean = torch.mean(x, dim = 1, keepdim = True)
         return (x - mean) / (var + self.eps).sqrt() * self.g + self.b
 
-class PreNorm(nn.Module):
-    def __init__(self, dim, fn):
-        super().__init__()
-        self.norm = LayerNorm(dim)
-        self.fn = fn
-    def forward(self, x, **kwargs):
-        x = self.norm(x)
-        return self.fn(x, **kwargs)
-
 class FeedForward(nn.Module):
     def __init__(self, dim, mult = 4, dropout = 0.):
         super().__init__()
         self.net = nn.Sequential(
+            LayerNorm(dim),
             nn.Conv2d(dim, dim * mult, 1),
             nn.GELU(),
             nn.Dropout(dropout),
@@ -75,6 +67,7 @@ class Attention(nn.Module):
         self.heads = heads
         self.scale = dim_head ** -0.5
 
+        self.norm = LayerNorm(dim)
         self.attend = nn.Softmax(dim = -1)
         self.dropout = nn.Dropout(dropout)
 
@@ -89,6 +82,8 @@ class Attention(nn.Module):
     def forward(self, x):
         shape = x.shape
         b, n, _, y, h = *shape, self.heads
+
+        x = self.norm(x)
         q, k, v = (self.to_q(x), *self.to_kv(x).chunk(2, dim = 1))
         q, k, v = map(lambda t: rearrange(t, 'b (h d) x y -> (b h) (x y) d', h = h), (q, k, v))
 
@@ -107,8 +102,8 @@ class Transformer(nn.Module):
         self.layers = nn.ModuleList([])
         for _ in range(depth):
             self.layers.append(nn.ModuleList([
-                PreNorm(dim, Attention(dim, proj_kernel = proj_kernel, kv_proj_stride = kv_proj_stride, heads = heads, dim_head = dim_head, dropout = dropout)),
-                PreNorm(dim, FeedForward(dim, mlp_mult, dropout = dropout))
+                Attention(dim, proj_kernel = proj_kernel, kv_proj_stride = kv_proj_stride, heads = heads, dim_head = dim_head, dropout = dropout),
+                FeedForward(dim, mlp_mult, dropout = dropout)
             ]))
     def forward(self, x):
         for attn, ff in self.layers:
@@ -145,12 +140,13 @@ class CvT(nn.Module):
         s3_heads = 6,
         s3_depth = 10,
         s3_mlp_mult = 4,
-        dropout = 0.
+        dropout = 0.,
+        channels = 3
     ):
         super().__init__()
         kwargs = dict(locals())
 
-        dim = 3
+        dim = channels
         layers = []
 
         for prefix in ('s1', 's2', 's3'):

vit_pytorch/deepvit.py

@@ -5,25 +5,11 @@ import torch.nn.functional as F
 from einops import rearrange, repeat
 from einops.layers.torch import Rearrange
 
-class Residual(nn.Module):
-    def __init__(self, fn):
-        super().__init__()
-        self.fn = fn
-    def forward(self, x, **kwargs):
-        return self.fn(x, **kwargs) + x
-
-class PreNorm(nn.Module):
-    def __init__(self, dim, fn):
-        super().__init__()
-        self.norm = nn.LayerNorm(dim)
-        self.fn = fn
-    def forward(self, x, **kwargs):
-        return self.fn(self.norm(x), **kwargs)
-
 class FeedForward(nn.Module):
     def __init__(self, dim, hidden_dim, dropout = 0.):
         super().__init__()
         self.net = nn.Sequential(
+            nn.LayerNorm(dim),
             nn.Linear(dim, hidden_dim),
             nn.GELU(),
             nn.Dropout(dropout),
@@ -40,6 +26,7 @@ class Attention(nn.Module):
         self.heads = heads
         self.scale = dim_head ** -0.5
 
+        self.norm = nn.LayerNorm(dim)
         self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
 
         self.dropout = nn.Dropout(dropout)
@@ -59,6 +46,8 @@ class Attention(nn.Module):
 
     def forward(self, x):
         b, n, _, h = *x.shape, self.heads
+        x = self.norm(x)
+
         qkv = self.to_qkv(x).chunk(3, dim = -1)
         q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), qkv)
 
@@ -86,13 +75,13 @@ class Transformer(nn.Module):
         self.layers = nn.ModuleList([])
         for _ in range(depth):
             self.layers.append(nn.ModuleList([
-                Residual(PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout))),
-                Residual(PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout)))
+                Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout),
+                FeedForward(dim, mlp_dim, dropout = dropout)
             ]))
     def forward(self, x):
         for attn, ff in self.layers:
-            x = attn(x)
-            x = ff(x)
+            x = attn(x) + x
+            x = ff(x) + x
         return x
 
 class DeepViT(nn.Module):

vit_pytorch/distill.py

@@ -1,6 +1,8 @@
 import torch
-import torch.nn.functional as F
 from torch import nn
+from torch.nn import Module
+import torch.nn.functional as F
+
 from vit_pytorch.vit import ViT
 from vit_pytorch.t2t import T2TViT
 from vit_pytorch.efficient import ViT as EfficientViT
@@ -12,6 +14,9 @@ from einops import rearrange, repeat
 def exists(val):
     return val is not None
 
+def default(val, d):
+    return val if exists(val) else d
+
 # classes
 
 class DistillMixin:
@@ -20,12 +25,12 @@ class DistillMixin:
         x = self.to_patch_embedding(img)
         b, n, _ = x.shape
 
-        cls_tokens = repeat(self.cls_token, '() n d -> b n d', b = b)
+        cls_tokens = repeat(self.cls_token, 'n d -> b n d', b = b)
         x = torch.cat((cls_tokens, x), dim = 1)
-        x += self.pos_embedding[:, :(n + 1)]
+        x += self.pos_embedding[:(n + 1)]
 
         if distilling:
-            distill_tokens = repeat(distill_token, '() n d -> b n d', b = b)
+            distill_tokens = repeat(distill_token, 'n d -> b n d', b = b)
             x = torch.cat((x, distill_tokens), dim = 1)
 
         x = self._attend(x)
@@ -97,7 +102,7 @@ class DistillableEfficientViT(DistillMixin, EfficientViT):
 
 # knowledge distillation wrapper
 
-class DistillWrapper(nn.Module):
+class DistillWrapper(Module):
     def __init__(
         self,
         *,
@@ -105,7 +110,8 @@ class DistillWrapper(nn.Module):
         student,
         temperature = 1.,
         alpha = 0.5,
-        hard = False
+        hard = False,
+        mlp_layernorm = False
     ):
         super().__init__()
         assert (isinstance(student, (DistillableViT, DistillableT2TViT, DistillableEfficientViT))) , 'student must be a vision transformer'
@@ -119,17 +125,17 @@ class DistillWrapper(nn.Module):
         self.alpha = alpha
         self.hard = hard
 
-        self.distillation_token = nn.Parameter(torch.randn(1, 1, dim))
+        self.distillation_token = nn.Parameter(torch.randn(1, dim))
 
         self.distill_mlp = nn.Sequential(
-            nn.LayerNorm(dim),
+            nn.LayerNorm(dim) if mlp_layernorm else nn.Identity(),
             nn.Linear(dim, num_classes)
         )
 
     def forward(self, img, labels, temperature = None, alpha = None, **kwargs):
-        b, *_ = img.shape
-        alpha = alpha if exists(alpha) else self.alpha
-        T = temperature if exists(temperature) else self.temperature
+
+        alpha = default(alpha, self.alpha)
+        T = default(temperature, self.temperature)
 
         with torch.no_grad():
             teacher_logits = self.teacher(img)

vit_pytorch/jet_vit.py

@@ -0,0 +1,359 @@
+import math
+import random
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.nn import Module, ModuleList
+
+from einops import rearrange, reduce, einsum
+from einops.layers.torch import Rearrange
+
+
+
+# helpers
+
+def pair(t):
+    return t if isinstance(t, tuple) else (t, t)
+
+def exists(val):
+    return val is not None
+
+# classes
+
+def linear_attn(q, k, v):
+    q = F.relu(q)
+    k = F.relu(k)
+
+    context = einsum(k, v, 'b h n d, b h n e -> b h d e')
+    normalizer = einsum(q, k.sum(dim=2), 'b h n d, b h d -> b h n')
+    attn = einsum(q, context, 'b h n d, b h d e -> b h n e')
+    return (attn / normalizer.unsqueeze(-1).clamp(min=1e-6))
+
+
+class FeedForward(Module):
+    def __init__(self, dim, hidden_dim, dropout = 0.):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Linear(dim, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, dim),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+class SqueezeDynamicConv(Module):
+    def __init__(self, dim, h_s, w_s, kernel_size = 3):
+        super().__init__()
+        self.kernel_size = kernel_size
+        self.padding = kernel_size // 2
+        self.dim = dim
+        self.h_s = h_s
+        self.w_s = w_s
+
+        self.mlp = nn.Sequential(
+            nn.Linear(dim, dim // 4),
+            nn.SiLU(),
+            nn.Linear(dim // 4, dim * kernel_size * kernel_size)
+        )
+
+    def forward(self, v):
+        b, heads, _, _ = v.shape
+
+        # Squeeze global context
+        v_mean = reduce(v, 'b h n d -> b (h d)', 'mean')
+
+        weight = self.mlp(v_mean)
+
+        weight = rearrange(weight, 'b (c k1 k2) -> (b c) 1 k1 k2', c = self.dim, k1 = self.kernel_size, k2 = self.kernel_size)
+
+        v_spatial = rearrange(v, 'b h (h_s w_s) d -> 1 (b h d) h_s w_s', h_s = self.h_s, w_s = self.w_s)
+
+        out = F.conv2d(v_spatial, weight, padding = self.padding, groups = b * self.dim)
+
+        return rearrange(out, '1 (b h d) h_s w_s -> b h (h_s w_s) d', b = b, h = heads, h_s = self.h_s, w_s = self.w_s)
+
+class WindowAttention(nn.Module):
+    def __init__(
+        self,
+        dim,
+        h_s,
+        w_s,
+        dim_head = 64,
+        dropout = 0.,
+        window_size = 7
+    ):
+        super().__init__()
+        assert (dim % dim_head) == 0, 'dimension should be divisible by dimension per head'
+
+        self.h_s = h_s
+        self.w_s = w_s
+        self.window_size = window_size
+
+        self.heads = dim // dim_head
+        self.scale = dim_head ** -0.5
+
+        self.norm = nn.LayerNorm(dim)
+        self.to_qkv = nn.Linear(dim, dim * 3, bias = False)
+
+        self.attend = nn.Sequential(
+            nn.Softmax(dim = -1),
+            nn.Dropout(dropout)
+        )
+
+        self.to_out = nn.Sequential(
+            nn.Linear(dim, dim, bias = False),
+            nn.Dropout(dropout)
+        )
+
+        self.rel_pos_bias = nn.Embedding((2 * window_size - 1) ** 2, self.heads)
+
+        pos = torch.arange(window_size)
+        grid = torch.stack(torch.meshgrid(pos, pos, indexing = 'ij'))
+        grid = rearrange(grid, 'c i j -> (i j) c')
+        rel_pos = rearrange(grid, 'i ... -> i 1 ...') - rearrange(grid, 'j ... -> 1 j ...')
+        rel_pos += window_size - 1
+        rel_pos_indices = (rel_pos * torch.tensor([2 * window_size - 1, 1])).sum(dim = -1)
+
+        self.register_buffer('rel_pos_indices', rel_pos_indices, persistent = False)
+
+    def forward(self, x):
+
+        x = self.norm(x)
+
+        x = rearrange(x, 'b (h w) d -> b h w d', h=self.h_s, w=self.w_s)
+        x = rearrange(x, 'b (x w1) (y w2) d -> b x y w1 w2 d',
+                      w1=self.window_size, w2=self.window_size)
+
+        batch, height, width, window_height, window_width, _, device, h = *x.shape, x.device, self.heads
+
+        x = rearrange(x, 'b x y w1 w2 d -> (b x y) (w1 w2) d')
+
+        q, k, v = self.to_qkv(x).chunk(3, dim=-1)
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=h), (q, k, v))
+
+        q = q * self.scale
+        sim = einsum(q, k, 'b h i d, b h j d -> b h i j')
+        sim = sim + rearrange(self.rel_pos_bias(self.rel_pos_indices), 'i j h -> h i j')
+        attn = self.attend(sim)
+
+        out = einsum(attn, v, 'b h i j, b h j d -> b h i d')
+
+        out = rearrange(out, 'b h (w1 w2) d -> b w1 w2 (h d)', w1=window_height, w2=window_width)
+
+        out = self.to_out(out)
+
+        out = rearrange(out, '(b x y) ... -> b x y ...', x=height, y=width)
+
+        out = rearrange(out, 'b x y w1 w2 d -> b (x w1) (y w2) d')
+
+        return rearrange(out, 'b h w d -> b (h w) d')
+
+
+class JetViTLinearAttention(Module):
+    def __init__(self, dim, h_s, w_s, heads = 8, dim_head = 64, dropout = 0., kernel_size = 3):
+        super().__init__()
+        inner_dim = dim_head * heads
+        project_out = not (heads == 1 and dim_head == dim)
+
+        self.heads = heads
+        self.h_s = h_s
+        self.w_s = w_s
+
+        self.norm = nn.LayerNorm(dim)
+        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
+
+        self.dynamic_conv = SqueezeDynamicConv(inner_dim, h_s, w_s, kernel_size = kernel_size)
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, dim),
+            nn.Dropout(dropout)
+        ) if project_out else nn.Identity()
+
+    def forward(self, x):
+        x = self.norm(x)
+
+        qkv = self.to_qkv(x).chunk(3, dim = -1)
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
+
+        linear_out = linear_attn(q, k, v)
+        linear_out = rearrange(linear_out, 'b h n d -> b n (h d)')
+
+        conv_out = self.dynamic_conv(v)
+        conv_out = rearrange(conv_out, 'b h n d -> b n (h d)')
+
+        return self.to_out(linear_out + conv_out)
+
+
+class Attention(Module):
+    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):
+        super().__init__()
+        inner_dim = dim_head *  heads
+        project_out = not (heads == 1 and dim_head == dim)
+
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+
+        self.norm = nn.LayerNorm(dim)
+
+        self.attend = nn.Softmax(dim = -1)
+        self.dropout = nn.Dropout(dropout)
+
+        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, dim),
+            nn.Dropout(dropout)
+        ) if project_out else nn.Identity()
+
+    def forward(self, x):
+        x = self.norm(x)
+
+        qkv = self.to_qkv(x).chunk(3, dim = -1)
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
+
+        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+
+        attn = self.attend(dots)
+        attn = self.dropout(attn)
+
+        out = torch.matmul(attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)
+
+
+
+class RandomLayerSelector(nn.Module):
+    def __init__(self, options: dict, attn_layer: list):
+        super().__init__()
+        self.options = nn.ModuleDict(options)
+        self.attn_layer = attn_layer
+
+    def forward(self, x):
+
+        if isinstance(self.attn_layer, tuple):
+            key = random.choice(self.attn_layer)
+        elif isinstance(self.attn_layer, str):
+            key = self.attn_layer
+
+        return self.options[key](x)
+
+
+class Transformer(nn.Module):
+    def __init__(
+        self,
+        dim,
+        depth,
+        heads,
+        dim_head,
+        mlp_dim,
+        h_s,
+        w_s,
+        dropout=0.,
+        window_size=7,
+        attn_layers=None,
+
+    ):
+        super().__init__()
+
+        if attn_layers is None:
+            attn_layers = ['FA'] * depth
+
+        self.norm = nn.LayerNorm(dim)
+        self.layers = nn.ModuleList([])
+
+        for i in range(depth):
+
+            attn = RandomLayerSelector(
+                options={
+                    'WA': WindowAttention(dim, h_s, w_s, dim_head=dim_head, dropout=dropout, window_size=window_size),
+                    'LA': JetViTLinearAttention(dim, h_s, w_s, heads=heads, dim_head=dim_head, dropout=dropout),
+                    'FA': Attention(dim, heads=heads, dim_head=dim_head, dropout=dropout) # Updated name to match your API
+                },
+                attn_layer=attn_layers[i]
+            )
+
+            self.layers.append(nn.ModuleList([
+                attn,
+                FeedForward(dim, mlp_dim, dropout=dropout)
+            ]))
+
+    def forward(self, x):
+        for attn, ff in self.layers:
+            x = attn(x) + x
+            x = ff(x) + x
+
+        return self.norm(x)
+
+
+class JetViT(Module):
+    def __init__(
+        self,
+        *,
+        image_size,
+        patch_size,
+        num_classes,
+        dim,
+        depth,
+        heads,
+        mlp_dim,
+        channels = 3,
+        dim_head = 64,
+        dropout = 0.,
+        emb_dropout = 0.,
+        window_size = 7,
+        attn_layers = None,
+    ):
+        super().__init__()
+        image_height, image_width = pair(image_size)
+        self.patch_size = patch_height, patch_width = pair(patch_size)
+
+        assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
+
+        h_s = image_height // patch_height
+        w_s = image_width // patch_width
+
+        num_patches = h_s * w_s
+        patch_dim = channels * patch_height * patch_width
+
+        self.to_patch_embedding = nn.Sequential(
+            Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_height, p2 = patch_width),
+            nn.LayerNorm(patch_dim),
+            nn.Linear(patch_dim, dim),
+            nn.LayerNorm(dim),
+        )
+
+        self.pos_embedding = nn.Parameter(torch.randn(num_patches, dim))
+        self.dropout = nn.Dropout(emb_dropout)
+
+        self.transformer = Transformer(
+            dim,
+            depth,
+            heads,
+            dim_head,
+            mlp_dim,
+            h_s,
+            w_s,
+            dropout,
+            window_size,
+            attn_layers,
+        )
+
+        self.to_latent = nn.Identity()
+        self.mlp_head = nn.Linear(dim, num_classes) if num_classes > 0 else None
+
+    def forward(self, img):
+        x = self.to_patch_embedding(img)
+        x = x + self.pos_embedding
+        x = self.dropout(x)
+
+        x = self.transformer(x)
+
+        if self.mlp_head is None:
+            return x
+
+        x = self.to_latent(x.mean(dim = 1))
+        return self.mlp_head(x)

vit_pytorch/jumbo_vit.py

@@ -0,0 +1,204 @@
+# Simpler Fast Vision Transformers with a Jumbo CLS Token
+# https://arxiv.org/abs/2502.15021
+
+import torch
+from torch import nn
+from torch.nn import Module, ModuleList
+
+from einops import rearrange, repeat, reduce, pack, unpack
+from einops.layers.torch import Rearrange
+
+# helpers
+
+def pair(t):
+    return t if isinstance(t, tuple) else (t, t)
+
+def divisible_by(num, den):
+    return (num % den) == 0
+
+def posemb_sincos_2d(h, w, dim, temperature: int = 10000, dtype = torch.float32):
+    y, x = torch.meshgrid(torch.arange(h), torch.arange(w), indexing="ij")
+    assert divisible_by(dim, 4), "feature dimension must be multiple of 4 for sincos emb"
+
+    omega = torch.arange(dim // 4) / (dim // 4 - 1)
+    omega = temperature ** -omega
+
+    y = y.flatten()[:, None] * omega[None, :]
+    x = x.flatten()[:, None] * omega[None, :]
+    pos_emb = torch.cat((x.sin(), x.cos(), y.sin(), y.cos()), dim=1)
+
+    return pos_emb.type(dtype)
+
+# classes
+
+def FeedForward(dim, mult = 4.):
+    hidden_dim = int(dim * mult)
+    return nn.Sequential(
+        nn.LayerNorm(dim),
+        nn.Linear(dim, hidden_dim),
+        nn.GELU(),
+        nn.Linear(hidden_dim, dim),
+    )
+
+class Attention(Module):
+    def __init__(self, dim, heads = 8, dim_head = 64):
+        super().__init__()
+        inner_dim = dim_head *  heads
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+        self.norm = nn.LayerNorm(dim)
+
+        self.attend = nn.Softmax(dim = -1)
+
+        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
+        self.to_out = nn.Linear(inner_dim, dim, bias = False)
+
+    def forward(self, x):
+        x = self.norm(x)
+
+        qkv = self.to_qkv(x).chunk(3, dim = -1)
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
+
+        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+
+        attn = self.attend(dots)
+
+        out = torch.matmul(attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)
+
+class JumboViT(Module):
+    def __init__(
+        self,
+        *,
+        image_size,
+        patch_size,
+        num_classes,
+        dim,
+        depth,
+        heads,
+        mlp_dim,
+        num_jumbo_cls = 1,  # differing from paper, allow for multiple jumbo cls, so one could break it up into 2 jumbo cls tokens with 3x the dim, as an example
+        jumbo_cls_k = 6,    # they use a CLS token with this factor times the dimension - 6 was the value they settled on
+        jumbo_ff_mult = 2,  # expansion factor of the jumbo cls token feedforward
+        channels = 3,
+        dim_head = 64
+    ):
+        super().__init__()
+        image_height, image_width = pair(image_size)
+        self.patch_size = patch_height, patch_width = pair(patch_size)
+
+        assert divisible_by(image_height, patch_height) and divisible_by(image_width, patch_width), 'Image dimensions must be divisible by the patch size.'
+
+        patch_dim = channels * patch_height * patch_width
+
+        self.to_patch_embedding = nn.Sequential(
+            Rearrange("b c (h p1) (w p2) -> b (h w) (p1 p2 c)", p1 = patch_height, p2 = patch_width),
+            nn.LayerNorm(patch_dim),
+            nn.Linear(patch_dim, dim),
+            nn.LayerNorm(dim),
+        )
+
+        self.pos_embedding = posemb_sincos_2d(
+            h = image_height // patch_height,
+            w = image_width // patch_width,
+            dim = dim,
+        )
+
+        jumbo_cls_dim = dim * jumbo_cls_k
+
+        self.jumbo_cls_token = nn.Parameter(torch.zeros(num_jumbo_cls, jumbo_cls_dim))
+
+        jumbo_cls_to_tokens = Rearrange('b n (k d) -> b (n k) d', k = jumbo_cls_k)
+        self.jumbo_cls_to_tokens = jumbo_cls_to_tokens
+
+        self.norm = nn.LayerNorm(dim)
+        self.layers = ModuleList([])
+
+        # attention and feedforwards
+
+        self.jumbo_ff = nn.Sequential(
+            Rearrange('b (n k) d -> b n (k d)', k = jumbo_cls_k),
+            FeedForward(jumbo_cls_dim, int(jumbo_cls_dim * jumbo_ff_mult)), # they use separate parameters for the jumbo feedforward, weight tied for parameter efficient
+            jumbo_cls_to_tokens
+        )
+
+        for _ in range(depth):
+            self.layers.append(ModuleList([
+                Attention(dim, heads = heads, dim_head = dim_head),
+                FeedForward(dim, mlp_dim),
+            ]))
+
+        self.to_latent = nn.Identity()
+
+        self.linear_head = nn.Linear(dim, num_classes)
+
+    def forward(self, img):
+
+        batch, device = img.shape[0], img.device
+
+        x = self.to_patch_embedding(img)
+
+        # pos embedding
+
+        pos_emb = self.pos_embedding.to(device, dtype = x.dtype)
+
+        x = x + pos_emb
+
+        # add cls tokens
+
+        cls_tokens = repeat(self.jumbo_cls_token, 'nj d -> b nj d', b = batch)
+
+        jumbo_tokens = self.jumbo_cls_to_tokens(cls_tokens)
+
+        x, cls_packed_shape = pack([jumbo_tokens, x], 'b * d')
+
+        # attention and feedforwards
+
+        for layer, (attn, ff) in enumerate(self.layers, start = 1):
+            is_last = layer == len(self.layers)
+
+            x = attn(x) + x
+
+            # jumbo feedforward
+
+            jumbo_cls_tokens, x = unpack(x, cls_packed_shape, 'b * d')
+
+            x = ff(x) + x
+            jumbo_cls_tokens = self.jumbo_ff(jumbo_cls_tokens) + jumbo_cls_tokens
+
+            if is_last:
+                continue
+
+            x, _ = pack([jumbo_cls_tokens, x], 'b * d')
+
+        pooled = reduce(jumbo_cls_tokens, 'b n d -> b d', 'mean')
+
+        # normalization and project to logits
+
+        embed = self.norm(pooled)
+
+        embed = self.to_latent(embed)
+        logits = self.linear_head(embed)
+        return logits
+
+# copy pasteable file
+
+if __name__ == '__main__':
+
+    v = JumboViT(
+        num_classes = 1000,
+        image_size = 64,
+        patch_size = 8,
+        dim = 16,
+        depth = 2,
+        heads = 2,
+        mlp_dim = 32,
+        jumbo_cls_k = 3,
+        jumbo_ff_mult = 2,
+    )
+
+    images = torch.randn(1, 3, 64, 64)
+
+    logits = v(images)
+    assert logits.shape == (1, 1000)

vit_pytorch/learnable_memory_vit.py

@@ -108,7 +108,7 @@ class ViT(nn.Module):
     def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, pool = 'cls', channels = 3, dim_head = 64, dropout = 0., emb_dropout = 0.):
         super().__init__()
         image_height, image_width = pair(image_size)
-        patch_height, patch_width = pair(patch_size)
+        self.patch_size = patch_height, patch_width = pair(patch_size)
 
         assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
 
@@ -145,7 +145,7 @@ class ViT(nn.Module):
         return x
 
     def forward(self, img):
-        x = self.img_to_tokens(img)        
+        x = self.img_to_tokens(img)
 
         x = self.transformer(x)
 
@@ -160,7 +160,7 @@ class Adapter(nn.Module):
         *,
         vit,
         num_memories_per_layer = 10,
-        num_classes = 2,   
+        num_classes = 2,
     ):
         super().__init__()
         assert isinstance(vit, ViT)
@@ -188,7 +188,7 @@ class Adapter(nn.Module):
         )
 
         # specialized attention mask to preserve the output of the original ViT
-        # it allows the memory CLS token to attend to all other tokens (and the learnable memory layer tokens), but not vice versa        
+        # it allows the memory CLS token to attend to all other tokens (and the learnable memory layer tokens), but not vice versa
 
         attn_mask = torch.ones((num_patches, num_patches), dtype = torch.bool)
         attn_mask = F.pad(attn_mask, (1, num_memories_per_layer), value = False)  # main tokens cannot attend to learnable memories per layer
@@ -203,7 +203,7 @@ class Adapter(nn.Module):
         # add task specific memory tokens
 
         memory_cls_tokens = repeat(self.memory_cls_token, 'd -> b 1 d', b = b)
-        tokens = torch.cat((memory_cls_tokens, tokens), dim = 1)        
+        tokens = torch.cat((memory_cls_tokens, tokens), dim = 1)
 
         # pass memories along with image tokens through transformer for attending
 

vit_pytorch/lejepa.py

@@ -0,0 +1,319 @@
+import random
+from functools import wraps
+
+import torch
+from torch import nn
+from torch.nn import Module
+import torch.nn.functional as F
+
+from torchvision import transforms as T
+from einops import rearrange
+
+# helper functions
+
+def exists(val):
+    return val is not None
+
+def default(val, d):
+    return val if exists(val) else d
+
+def singleton(cache_key):
+    def inner_fn(fn):
+        @wraps(fn)
+        def wrapper(self, *args, **kwargs):
+            instance = getattr(self, cache_key)
+            if instance is not None:
+                return instance
+
+            instance = fn(self, *args, **kwargs)
+            setattr(self, cache_key, instance)
+            return instance
+        return wrapper
+    return inner_fn
+
+def get_module_device(module):
+    return next(module.parameters()).device
+
+def l2norm(t, eps = 1e-6):
+    return F.normalize(t, dim = -1, eps = eps)
+
+# loss function
+
+def sigreg_loss(
+    x,
+    num_slices = 1024,
+    domain = (-5, 5),
+    num_knots = 17
+):
+    # Randall Balestriero - https://arxiv.org/abs/2511.08544
+
+    dim, device = x.shape[-1], x.device
+
+    # slice sampling
+
+    rand_projs = torch.randn((num_slices, dim), device = device)
+    rand_projs = l2norm(rand_projs)
+
+    # integration points
+
+    t = torch.linspace(*domain, num_knots, device = device)
+
+    # theoretical CF for N(0, 1) and Gauss. window
+
+    exp_f = (-0.5 * t.square()).exp()
+
+    # empirical CF
+
+    x_t = torch.einsum('... d, m d -> ... m', x, rand_projs)
+    x_t = rearrange(x_t, '... m -> (...) m')
+
+    x_t = rearrange(x_t, 'n m -> n m 1') * t
+    ecf = (1j * x_t).exp().mean(dim = 0)
+
+    # weighted L2 distance
+
+    err = ecf.sub(exp_f).abs().square().mul(exp_f)
+
+    return torch.trapezoid(err, t, dim = -1).mean()
+
+# augmentation utils
+
+class RandomApply(Module):
+    def __init__(self, fn, p):
+        super().__init__()
+        self.fn = fn
+        self.p = p
+
+    def forward(self, x):
+        if random.random() > self.p:
+            return x
+        return self.fn(x)
+
+# MLP class for projector
+
+class L2Norm(Module):
+    def forward(self, x, eps = 1e-6):
+        return l2norm(x, eps)
+
+class MLP(Module):
+    def __init__(self, dim, dim_out, num_layers, hidden_size = 256):
+        super().__init__()
+
+        layers = []
+        dims = (dim, *((hidden_size,) * (num_layers - 1)))
+
+        for ind, (layer_dim_in, layer_dim_out) in enumerate(zip(dims[:-1], dims[1:])):
+            is_last = ind == (len(dims) - 1)
+
+            layers.extend([
+                nn.Linear(layer_dim_in, layer_dim_out),
+                nn.GELU() if not is_last else nn.Identity()
+            ])
+
+        self.net = nn.Sequential(
+            *layers,
+            L2Norm(),
+            nn.Linear(hidden_size, dim_out)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+# wrapper
+
+class NetWrapper(Module):
+    def __init__(self, net, output_dim, projection_hidden_size, projection_num_layers, layer = -2):
+        super().__init__()
+        self.net = net
+        self.layer = layer
+
+        self.projector = None
+        self.projection_hidden_size = projection_hidden_size
+        self.projection_num_layers = projection_num_layers
+        self.output_dim = output_dim
+
+        self.hidden = {}
+        self.hook_registered = False
+
+    def _find_layer(self):
+        if type(self.layer) == str:
+            modules = dict([*self.net.named_modules()])
+            return modules.get(self.layer, None)
+        elif type(self.layer) == int:
+            children = [*self.net.children()]
+            return children[self.layer]
+        return None
+
+    def _hook(self, _, input, output):
+        device = input[0].device
+        self.hidden[device] = output.flatten(1)
+
+    def _register_hook(self):
+        layer = self._find_layer()
+        assert layer is not None, f'hidden layer ({self.layer}) not found'
+        handle = layer.register_forward_hook(self._hook)
+        self.hook_registered = True
+
+    @singleton('projector')
+    def _get_projector(self, hidden):
+        _, dim = hidden.shape
+        projector = MLP(dim, self.output_dim, self.projection_num_layers, self.projection_hidden_size)
+        return projector.to(hidden)
+
+    def get_embedding(self, x):
+        if self.layer == -1:
+            return self.net(x)
+
+        if not self.hook_registered:
+            self._register_hook()
+
+        self.hidden.clear()
+        _ = self.net(x)
+        hidden = self.hidden[x.device]
+        self.hidden.clear()
+
+        assert hidden is not None, f'hidden layer {self.layer} never emitted an output'
+        return hidden
+
+    def forward(self, x, return_projection = True):
+        embed = self.get_embedding(x)
+        if not return_projection:
+            return embed
+
+        projector = self._get_projector(embed)
+        return projector(embed), embed
+
+# main class
+
+class LeJEPA(nn.Module):
+    def __init__(
+        self,
+        net,
+        image_size,
+        hidden_layer = -2,
+        projection_hidden_size = 256,
+        num_classes_K = 65336,
+        projection_layers = 4,
+        local_upper_crop_scale = 0.4,
+        global_lower_crop_scale = 0.5,
+        target_loss_weight = 1.,
+        sigreg_loss_weight = 1.,
+        sigreg_loss_kwargs = dict(
+            num_slices = 1024,
+            domain = (-5, 5),
+            num_knots = 17
+        ),
+        augment_fn = None,
+        augment_fn2 = None
+    ):
+        super().__init__()
+        self.net = net
+
+        # default BYOL augmentation
+
+        DEFAULT_AUG = torch.nn.Sequential(
+            RandomApply(
+                T.ColorJitter(0.8, 0.8, 0.8, 0.2),
+                p = 0.3
+            ),
+            T.RandomGrayscale(p=0.2),
+            T.RandomHorizontalFlip(),
+            RandomApply(
+                T.GaussianBlur((3, 3), (1.0, 2.0)),
+                p = 0.2
+            ),
+            T.Normalize(
+                mean=torch.tensor([0.485, 0.456, 0.406]),
+                std=torch.tensor([0.229, 0.224, 0.225])),
+        )
+
+        self.augment1 = default(augment_fn, DEFAULT_AUG)
+        self.augment2 = default(augment_fn2, DEFAULT_AUG)
+
+        # local and global crops
+
+        self.local_crop = T.RandomResizedCrop((image_size, image_size), scale = (0.05, local_upper_crop_scale))
+        self.global_crop = T.RandomResizedCrop((image_size, image_size), scale = (global_lower_crop_scale, 1.))
+
+        self.encoder = NetWrapper(net, num_classes_K, projection_hidden_size, projection_layers, layer = hidden_layer)
+
+        self.target_loss_weight = target_loss_weight
+        self.sigreg_loss_weight = sigreg_loss_weight
+        self.sigreg_loss_kwargs = sigreg_loss_kwargs
+
+        # get device of network and make wrapper same device
+        device = get_module_device(net)
+        self.to(device)
+
+        # send a mock image tensor to instantiate singleton parameters
+        self.forward(torch.randn(2, 3, image_size, image_size, device=device))
+
+    def forward(
+        self,
+        x,
+        return_embedding = False,
+        return_projection = True
+    ):
+        if return_embedding:
+            return self.encoder(x, return_projection = return_projection)
+
+        image_one, image_two = self.augment1(x), self.augment2(x)
+
+        local_image_one, local_image_two   = self.local_crop(image_one),  self.local_crop(image_two)
+        global_image_one, global_image_two = self.global_crop(image_one), self.global_crop(image_two)
+
+        local_images = torch.cat((local_image_one, local_image_two), dim = 0)
+        proj_locals, _ = self.encoder(local_images)
+        proj_local_one, proj_local_two = proj_locals.chunk(2, dim = 0)
+
+        with torch.no_grad():
+            global_images = torch.cat((global_image_one, global_image_two), dim = 0)
+            proj_globals, _ = self.encoder(global_images)
+            proj_global_one, proj_global_two = proj_globals.chunk(2, dim = 0)
+
+        # invariance loss
+
+        mse_loss = F.mse_loss(proj_local_one, proj_global_two) + F.mse_loss(proj_local_two, proj_global_one)
+
+        # sigreg loss
+
+        sreg_loss = sigreg_loss(proj_locals, **self.sigreg_loss_kwargs)
+
+        return mse_loss * self.target_loss_weight + sreg_loss * self.sigreg_loss_weight
+
+# quick run
+
+if __name__ == '__main__':
+    from vit_pytorch import ViT
+
+    model = ViT(
+        image_size = 256,
+        patch_size = 32,
+        num_classes = 1000,
+        dim = 1024,
+        depth = 6,
+        heads = 8,
+        mlp_dim = 2048
+    )
+
+    learner = LeJEPA(
+        model,
+        image_size = 256,
+        hidden_layer = 'to_latent',          # layer name where output is hidden dimension
+        projection_hidden_size = 256,        # projector network hidden dimension
+        projection_layers = 4,               # number of layers in projection network
+        num_classes_K = 65336,               # output dimension
+        target_loss_weight = 1.0,
+        sigreg_loss_weight = 1.0
+    )
+
+    opt = torch.optim.Adam(learner.parameters(), lr = 3e-4)
+
+    images = torch.randn(8, 3, 256, 256)
+
+    loss = learner(images)
+    opt.zero_grad()
+    loss.backward()
+    opt.step()
+
+    print('loss:', loss.item())

vit_pytorch/levit.py

@@ -182,7 +182,7 @@ class LeViT(nn.Module):
     def forward(self, img):
         x = self.conv_embedding(img)
 
-        x = self.backbone(x)        
+        x = self.backbone(x)
 
         x = self.pool(x)
 

vit_pytorch/local_vit.py

@@ -26,16 +26,6 @@ class ExcludeCLS(nn.Module):
         x = self.fn(x, **kwargs)
         return torch.cat((cls_token, x), dim = 1)
 
-# prenorm
-
-class PreNorm(nn.Module):
-    def __init__(self, dim, fn):
-        super().__init__()
-        self.norm = nn.LayerNorm(dim)
-        self.fn = fn
-    def forward(self, x, **kwargs):
-        return self.fn(self.norm(x), **kwargs)
-
 # feed forward related classes
 
 class DepthWiseConv2d(nn.Module):
@@ -52,6 +42,7 @@ class FeedForward(nn.Module):
     def __init__(self, dim, hidden_dim, dropout = 0.):
         super().__init__()
         self.net = nn.Sequential(
+            nn.LayerNorm(dim),
             nn.Conv2d(dim, hidden_dim, 1),
             nn.Hardswish(),
             DepthWiseConv2d(hidden_dim, hidden_dim, 3, padding = 1),
@@ -77,6 +68,7 @@ class Attention(nn.Module):
         self.heads = heads
         self.scale = dim_head ** -0.5
 
+        self.norm = nn.LayerNorm(dim)
         self.attend = nn.Softmax(dim = -1)
         self.dropout = nn.Dropout(dropout)
         self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
@@ -88,6 +80,8 @@ class Attention(nn.Module):
 
     def forward(self, x):
         b, n, _, h = *x.shape, self.heads
+
+        x = self.norm(x)
         qkv = self.to_qkv(x).chunk(3, dim = -1)
         q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), qkv)
 
@@ -106,8 +100,8 @@ class Transformer(nn.Module):
         self.layers = nn.ModuleList([])
         for _ in range(depth):
             self.layers.append(nn.ModuleList([
-                Residual(PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout))),
-                ExcludeCLS(Residual(PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout))))
+                Residual(Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout)),
+                ExcludeCLS(Residual(FeedForward(dim, mlp_dim, dropout = dropout)))
             ]))
     def forward(self, x):
         for attn, ff in self.layers:

vit_pytorch/look_vit.py

@@ -0,0 +1,278 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.nn import Module, ModuleList
+
+from einops import einsum, rearrange, repeat, reduce
+from einops.layers.torch import Rearrange
+
+# helpers
+
+def exists(val):
+    return val is not None
+
+def default(val, d):
+    return val if exists(val) else d
+
+def divisible_by(num, den):
+    return (num % den) == 0
+
+# simple vit sinusoidal pos emb
+
+def posemb_sincos_2d(t, temperature = 10000):
+    h, w, d, device = *t.shape[1:], t.device
+    y, x = torch.meshgrid(torch.arange(h, device = device), torch.arange(w, device = device), indexing = 'ij')
+    assert (d % 4) == 0, "feature dimension must be multiple of 4 for sincos emb"
+    omega = torch.arange(d // 4, device = device) / (d // 4 - 1)
+    omega = temperature ** -omega
+
+    y = y.flatten()[:, None] * omega[None, :]
+    x = x.flatten()[:, None] * omega[None, :]
+    pos = torch.cat((x.sin(), x.cos(), y.sin(), y.cos()), dim = 1)
+
+    return pos.float()
+
+# bias-less layernorm with unit offset trick (discovered by Ohad Rubin)
+
+class LayerNorm(Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.ln = nn.LayerNorm(dim, elementwise_affine = False)
+        self.gamma = nn.Parameter(torch.zeros(dim))
+
+    def forward(self, x):
+        normed = self.ln(x)
+        return normed * (self.gamma + 1)
+
+# mlp
+
+def MLP(dim, factor = 4, dropout = 0.):
+    hidden_dim = int(dim * factor)
+    return nn.Sequential(
+        LayerNorm(dim),
+        nn.Linear(dim, hidden_dim),
+        nn.GELU(),
+        nn.Dropout(dropout),
+        nn.Linear(hidden_dim, dim),
+        nn.Dropout(dropout)
+    )
+
+# attention
+
+class Attention(Module):
+    def __init__(
+        self,
+        dim,
+        heads = 8,
+        dim_head = 64,
+        dropout = 0.,
+        cross_attend = False,
+        reuse_attention = False
+    ):
+        super().__init__()
+        inner_dim = dim_head *  heads
+
+        self.scale = dim_head ** -0.5
+        self.heads = heads
+        self.reuse_attention = reuse_attention
+        self.cross_attend = cross_attend
+
+        self.split_heads = Rearrange('b n (h d) -> b h n d', h = heads)
+
+        self.norm = LayerNorm(dim) if not reuse_attention else nn.Identity()
+        self.norm_context = LayerNorm(dim) if cross_attend else nn.Identity()
+
+        self.attend = nn.Softmax(dim = -1)
+        self.dropout = nn.Dropout(dropout)
+
+        self.to_q = nn.Linear(dim, inner_dim, bias = False) if not reuse_attention else None
+        self.to_k = nn.Linear(dim, inner_dim, bias = False) if not reuse_attention else None
+        self.to_v = nn.Linear(dim, inner_dim, bias = False)
+
+        self.to_out = nn.Sequential(
+            Rearrange('b h n d -> b n (h d)'),
+            nn.Linear(inner_dim, dim, bias = False),
+            nn.Dropout(dropout)
+        )
+
+    def forward(
+        self,
+        x,
+        context = None,
+        return_qk_sim = False,
+        qk_sim = None
+    ):
+        x = self.norm(x)
+
+        assert not (exists(context) ^ self.cross_attend)
+
+        if self.cross_attend:
+            context = self.norm_context(context)
+        else:
+            context = x
+
+        v = self.to_v(context)
+        v = self.split_heads(v)
+
+        if not self.reuse_attention:
+            qk = (self.to_q(x), self.to_k(context))
+            q, k = tuple(self.split_heads(t) for t in qk)
+
+            q = q * self.scale
+            qk_sim = einsum(q, k, 'b h i d, b h j d -> b h i j')
+
+        else:
+            assert exists(qk_sim), 'qk sim matrix must be passed in for reusing previous attention'
+
+        attn = self.attend(qk_sim)
+        attn = self.dropout(attn)
+
+        out = einsum(attn, v, 'b h i j, b h j d -> b h i d')
+        out = self.to_out(out)
+
+        if not return_qk_sim:
+            return out
+
+        return out, qk_sim
+
+# LookViT
+
+class LookViT(Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        image_size,
+        num_classes,
+        depth = 3,
+        patch_size = 16,
+        heads = 8,
+        mlp_factor = 4,
+        dim_head = 64,
+        highres_patch_size = 12,
+        highres_mlp_factor = 4,
+        cross_attn_heads = 8,
+        cross_attn_dim_head = 64,
+        patch_conv_kernel_size = 7,
+        dropout = 0.1,
+        channels = 3
+    ):
+        super().__init__()
+        assert divisible_by(image_size, highres_patch_size)
+        assert divisible_by(image_size, patch_size)
+        assert patch_size > highres_patch_size, 'patch size of the main vision transformer should be smaller than the highres patch sizes (that does the `lookup`)'
+        assert not divisible_by(patch_conv_kernel_size, 2)
+
+        self.dim = dim
+        self.image_size = image_size
+        self.patch_size = patch_size
+
+        kernel_size = patch_conv_kernel_size
+        patch_dim = (highres_patch_size * highres_patch_size) * channels
+
+        self.to_patches = nn.Sequential(
+            Rearrange('b c (h p1) (w p2) -> b (p1 p2 c) h w', p1 = highres_patch_size, p2 = highres_patch_size),
+            nn.Conv2d(patch_dim, dim, kernel_size, padding = kernel_size // 2),
+            Rearrange('b c h w -> b h w c'),
+            LayerNorm(dim),
+        )
+
+        # absolute positions
+
+        num_patches = (image_size // highres_patch_size) ** 2
+        self.pos_embedding = nn.Parameter(torch.randn(num_patches, dim))
+
+        # lookvit blocks
+
+        layers = ModuleList([])
+
+        for _ in range(depth):
+            layers.append(ModuleList([
+                Attention(dim = dim, dim_head = dim_head, heads = heads, dropout = dropout),
+                MLP(dim = dim, factor = mlp_factor, dropout = dropout),
+                Attention(dim = dim, dim_head = cross_attn_dim_head, heads = cross_attn_heads, dropout = dropout, cross_attend = True),
+                Attention(dim = dim, dim_head = cross_attn_dim_head, heads = cross_attn_heads, dropout = dropout, cross_attend = True, reuse_attention = True),
+                LayerNorm(dim),
+                MLP(dim = dim, factor = highres_mlp_factor, dropout = dropout)
+            ]))
+
+        self.layers = layers
+
+        self.norm = LayerNorm(dim)
+        self.highres_norm = LayerNorm(dim)
+
+        self.to_logits = nn.Linear(dim, num_classes, bias = False)
+
+    def forward(self, img):
+        assert img.shape[-2:] == (self.image_size, self.image_size)
+
+        # to patch tokens and positions
+
+        highres_tokens = self.to_patches(img)
+        size = highres_tokens.shape[-2]
+
+        pos_emb = posemb_sincos_2d(highres_tokens)
+        highres_tokens = highres_tokens + rearrange(pos_emb, '(h w) d -> h w d', h = size)
+
+        tokens = F.interpolate(
+            rearrange(highres_tokens, 'b h w d -> b d h w'),
+            img.shape[-1] // self.patch_size,
+            mode = 'bilinear'
+        )
+
+        tokens = rearrange(tokens, 'b c h w -> b (h w) c')
+        highres_tokens = rearrange(highres_tokens, 'b h w c -> b (h w) c')
+
+        # attention and feedforwards
+
+        for attn, mlp, lookup_cross_attn, highres_attn, highres_norm, highres_mlp in self.layers:
+
+            # main tokens cross attends (lookup) on the high res tokens
+
+            lookup_out, qk_sim = lookup_cross_attn(tokens, highres_tokens, return_qk_sim = True)  # return attention as they reuse the attention matrix
+            tokens = lookup_out + tokens
+
+            tokens = attn(tokens) + tokens
+            tokens = mlp(tokens) + tokens
+
+            # attention-reuse
+
+            qk_sim = rearrange(qk_sim, 'b h i j -> b h j i') # transpose for reverse cross attention
+
+            highres_tokens = highres_attn(highres_tokens, tokens, qk_sim = qk_sim) + highres_tokens
+            highres_tokens = highres_norm(highres_tokens)
+
+            highres_tokens = highres_mlp(highres_tokens) + highres_tokens
+
+        # to logits
+
+        tokens = self.norm(tokens)
+        highres_tokens = self.highres_norm(highres_tokens)
+
+        tokens = reduce(tokens, 'b n d -> b d', 'mean')
+        highres_tokens = reduce(highres_tokens, 'b n d -> b d', 'mean')
+
+        return self.to_logits(tokens + highres_tokens)
+
+# main
+
+if __name__ == '__main__':
+    v = LookViT(
+        image_size = 256,
+        num_classes = 1000,
+        dim = 512,
+        depth = 2,
+        heads = 8,
+        dim_head = 64,
+        patch_size = 32,
+        highres_patch_size = 8,
+        highres_mlp_factor = 2,
+        cross_attn_heads = 8,
+        cross_attn_dim_head = 64,
+        dropout = 0.1
+    ).cuda()
+
+    img = torch.randn(2, 3, 256, 256).cuda()
+    pred = v(img)
+
+    assert pred.shape == (2, 1000)

vit_pytorch/mae.py

@@ -52,7 +52,7 @@ class MAE(nn.Module):
         if self.encoder.pool == "cls":
             tokens += self.encoder.pos_embedding[:, 1:(num_patches + 1)]
         elif self.encoder.pool == "mean":
-            tokens += self.encoder.pos_embedding.to(device, dtype=tokens.dtype) 
+            tokens += self.encoder.pos_embedding.to(device, dtype=tokens.dtype)
 
         # calculate of patches needed to be masked, and get random indices, dividing it up for mask vs unmasked
 
@@ -87,7 +87,7 @@ class MAE(nn.Module):
         mask_tokens = mask_tokens + self.decoder_pos_emb(masked_indices)
 
         # concat the masked tokens to the decoder tokens and attend with decoder
-        
+
         decoder_tokens = torch.zeros(batch, num_patches, self.decoder_dim, device=device)
         decoder_tokens[batch_range, unmasked_indices] = unmasked_decoder_tokens
         decoder_tokens[batch_range, masked_indices] = mask_tokens

vit_pytorch/max_vit.py

@@ -19,20 +19,20 @@ def cast_tuple(val, length = 1):
 
 # helper classes
 
-class PreNormResidual(nn.Module):
+class Residual(nn.Module):
     def __init__(self, dim, fn):
         super().__init__()
-        self.norm = nn.LayerNorm(dim)
         self.fn = fn
 
     def forward(self, x):
-        return self.fn(self.norm(x)) + x
+        return self.fn(x) + x
 
 class FeedForward(nn.Module):
     def __init__(self, dim, mult = 4, dropout = 0.):
         super().__init__()
         inner_dim = int(dim * mult)
         self.net = nn.Sequential(
+            nn.LayerNorm(dim),
             nn.Linear(dim, inner_dim),
             nn.GELU(),
             nn.Dropout(dropout),
@@ -77,7 +77,7 @@ class Dropsample(nn.Module):
     def __init__(self, prob = 0):
         super().__init__()
         self.prob = prob
-  
+
     def forward(self, x):
         device = x.device
 
@@ -132,6 +132,7 @@ class Attention(nn.Module):
         self.heads = dim // dim_head
         self.scale = dim_head ** -0.5
 
+        self.norm = nn.LayerNorm(dim)
         self.to_qkv = nn.Linear(dim, dim * 3, bias = False)
 
         self.attend = nn.Sequential(
@@ -160,6 +161,8 @@ class Attention(nn.Module):
     def forward(self, x):
         batch, height, width, window_height, window_width, _, device, h = *x.shape, x.device, self.heads
 
+        x = self.norm(x)
+
         # flatten
 
         x = rearrange(x, 'b x y w1 w2 d -> (b x y) (w1 w2) d')
@@ -170,7 +173,7 @@ class Attention(nn.Module):
 
         # split heads
 
-        q, k, v = map(lambda t: rearrange(t, 'b n (h d ) -> b h n d', h = h), (q, k, v))
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), (q, k, v))
 
         # scale
 
@@ -259,13 +262,13 @@ class MaxViT(nn.Module):
                         shrinkage_rate = mbconv_shrinkage_rate
                     ),
                     Rearrange('b d (x w1) (y w2) -> b x y w1 w2 d', w1 = w, w2 = w),  # block-like attention
-                    PreNormResidual(layer_dim, Attention(dim = layer_dim, dim_head = dim_head, dropout = dropout, window_size = w)),
-                    PreNormResidual(layer_dim, FeedForward(dim = layer_dim, dropout = dropout)),
+                    Residual(layer_dim, Attention(dim = layer_dim, dim_head = dim_head, dropout = dropout, window_size = w)),
+                    Residual(layer_dim, FeedForward(dim = layer_dim, dropout = dropout)),
                     Rearrange('b x y w1 w2 d -> b d (x w1) (y w2)'),
 
                     Rearrange('b d (w1 x) (w2 y) -> b x y w1 w2 d', w1 = w, w2 = w),  # grid-like attention
-                    PreNormResidual(layer_dim, Attention(dim = layer_dim, dim_head = dim_head, dropout = dropout, window_size = w)),
-                    PreNormResidual(layer_dim, FeedForward(dim = layer_dim, dropout = dropout)),
+                    Residual(layer_dim, Attention(dim = layer_dim, dim_head = dim_head, dropout = dropout, window_size = w)),
+                    Residual(layer_dim, FeedForward(dim = layer_dim, dropout = dropout)),
                     Rearrange('b x y w1 w2 d -> b d (w1 x) (w2 y)'),
                 )
 

vit_pytorch/max_vit_with_registers.py

@@ -0,0 +1,340 @@
+from functools import partial
+
+import torch
+from torch import nn, einsum
+import torch.nn.functional as F
+from torch.nn import Module, ModuleList, Sequential
+
+from einops import rearrange, repeat, reduce, pack, unpack
+from einops.layers.torch import Rearrange, Reduce
+
+# helpers
+
+def exists(val):
+    return val is not None
+
+def default(val, d):
+    return val if exists(val) else d
+
+def pack_one(x, pattern):
+    return pack([x], pattern)
+
+def unpack_one(x, ps, pattern):
+    return unpack(x, ps, pattern)[0]
+
+def cast_tuple(val, length = 1):
+    return val if isinstance(val, tuple) else ((val,) * length)
+
+# helper classes
+
+def FeedForward(dim, mult = 4, dropout = 0.):
+    inner_dim = int(dim * mult)
+    return Sequential(
+        nn.LayerNorm(dim),
+        nn.Linear(dim, inner_dim),
+        nn.GELU(),
+        nn.Dropout(dropout),
+        nn.Linear(inner_dim, dim),
+        nn.Dropout(dropout)
+    )
+
+# MBConv
+
+class SqueezeExcitation(Module):
+    def __init__(self, dim, shrinkage_rate = 0.25):
+        super().__init__()
+        hidden_dim = int(dim * shrinkage_rate)
+
+        self.gate = Sequential(
+            Reduce('b c h w -> b c', 'mean'),
+            nn.Linear(dim, hidden_dim, bias = False),
+            nn.SiLU(),
+            nn.Linear(hidden_dim, dim, bias = False),
+            nn.Sigmoid(),
+            Rearrange('b c -> b c 1 1')
+        )
+
+    def forward(self, x):
+        return x * self.gate(x)
+
+class MBConvResidual(Module):
+    def __init__(self, fn, dropout = 0.):
+        super().__init__()
+        self.fn = fn
+        self.dropsample = Dropsample(dropout)
+
+    def forward(self, x):
+        out = self.fn(x)
+        out = self.dropsample(out)
+        return out + x
+
+class Dropsample(Module):
+    def __init__(self, prob = 0):
+        super().__init__()
+        self.prob = prob
+
+    def forward(self, x):
+        device = x.device
+
+        if self.prob == 0. or (not self.training):
+            return x
+
+        keep_mask = torch.FloatTensor((x.shape[0], 1, 1, 1), device = device).uniform_() > self.prob
+        return x * keep_mask / (1 - self.prob)
+
+def MBConv(
+    dim_in,
+    dim_out,
+    *,
+    downsample,
+    expansion_rate = 4,
+    shrinkage_rate = 0.25,
+    dropout = 0.
+):
+    hidden_dim = int(expansion_rate * dim_out)
+    stride = 2 if downsample else 1
+
+    net = Sequential(
+        nn.Conv2d(dim_in, hidden_dim, 1),
+        nn.BatchNorm2d(hidden_dim),
+        nn.GELU(),
+        nn.Conv2d(hidden_dim, hidden_dim, 3, stride = stride, padding = 1, groups = hidden_dim),
+        nn.BatchNorm2d(hidden_dim),
+        nn.GELU(),
+        SqueezeExcitation(hidden_dim, shrinkage_rate = shrinkage_rate),
+        nn.Conv2d(hidden_dim, dim_out, 1),
+        nn.BatchNorm2d(dim_out)
+    )
+
+    if dim_in == dim_out and not downsample:
+        net = MBConvResidual(net, dropout = dropout)
+
+    return net
+
+# attention related classes
+
+class Attention(Module):
+    def __init__(
+        self,
+        dim,
+        dim_head = 32,
+        dropout = 0.,
+        window_size = 7,
+        num_registers = 1
+    ):
+        super().__init__()
+        assert num_registers > 0
+        assert (dim % dim_head) == 0, 'dimension should be divisible by dimension per head'
+
+        self.heads = dim // dim_head
+        self.scale = dim_head ** -0.5
+
+        self.norm = nn.LayerNorm(dim)
+        self.to_qkv = nn.Linear(dim, dim * 3, bias = False)
+
+        self.attend = nn.Sequential(
+            nn.Softmax(dim = -1),
+            nn.Dropout(dropout)
+        )
+
+        self.to_out = nn.Sequential(
+            nn.Linear(dim, dim, bias = False),
+            nn.Dropout(dropout)
+        )
+
+        # relative positional bias
+
+        num_rel_pos_bias = (2 * window_size - 1) ** 2
+
+        self.rel_pos_bias = nn.Embedding(num_rel_pos_bias + 1, self.heads)
+
+        pos = torch.arange(window_size)
+        grid = torch.stack(torch.meshgrid(pos, pos, indexing = 'ij'))
+        grid = rearrange(grid, 'c i j -> (i j) c')
+        rel_pos = rearrange(grid, 'i ... -> i 1 ...') - rearrange(grid, 'j ... -> 1 j ...')
+        rel_pos += window_size - 1
+        rel_pos_indices = (rel_pos * torch.tensor([2 * window_size - 1, 1])).sum(dim = -1)
+
+        rel_pos_indices = F.pad(rel_pos_indices, (num_registers, 0, num_registers, 0), value = num_rel_pos_bias)
+        self.register_buffer('rel_pos_indices', rel_pos_indices, persistent = False)
+
+    def forward(self, x):
+        device, h, bias_indices = x.device, self.heads, self.rel_pos_indices
+
+        x = self.norm(x)
+
+        # project for queries, keys, values
+
+        q, k, v = self.to_qkv(x).chunk(3, dim = -1)
+
+        # split heads
+
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), (q, k, v))
+
+        # scale
+
+        q = q * self.scale
+
+        # sim
+
+        sim = einsum('b h i d, b h j d -> b h i j', q, k)
+
+        # add positional bias
+
+        bias = self.rel_pos_bias(bias_indices)
+        sim = sim + rearrange(bias, 'i j h -> h i j')
+
+        # attention
+
+        attn = self.attend(sim)
+
+        # aggregate
+
+        out = einsum('b h i j, b h j d -> b h i d', attn, v)
+
+        # combine heads out
+
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)
+
+class MaxViT(Module):
+    def __init__(
+        self,
+        *,
+        num_classes,
+        dim,
+        depth,
+        dim_head = 32,
+        dim_conv_stem = None,
+        window_size = 7,
+        mbconv_expansion_rate = 4,
+        mbconv_shrinkage_rate = 0.25,
+        dropout = 0.1,
+        channels = 3,
+        num_register_tokens = 4
+    ):
+        super().__init__()
+        assert isinstance(depth, tuple), 'depth needs to be tuple if integers indicating number of transformer blocks at that stage'
+        assert num_register_tokens > 0
+
+        # convolutional stem
+
+        dim_conv_stem = default(dim_conv_stem, dim)
+
+        self.conv_stem = Sequential(
+            nn.Conv2d(channels, dim_conv_stem, 3, stride = 2, padding = 1),
+            nn.Conv2d(dim_conv_stem, dim_conv_stem, 3, padding = 1)
+        )
+
+        # variables
+
+        num_stages = len(depth)
+
+        dims = tuple(map(lambda i: (2 ** i) * dim, range(num_stages)))
+        dims = (dim_conv_stem, *dims)
+        dim_pairs = tuple(zip(dims[:-1], dims[1:]))
+
+        self.layers = nn.ModuleList([])
+
+        # window size
+
+        self.window_size = window_size
+
+        self.register_tokens = nn.ParameterList([])
+
+        # iterate through stages
+
+        for ind, ((layer_dim_in, layer_dim), layer_depth) in enumerate(zip(dim_pairs, depth)):
+            for stage_ind in range(layer_depth):
+                is_first = stage_ind == 0
+                stage_dim_in = layer_dim_in if is_first else layer_dim
+
+                conv = MBConv(
+                    stage_dim_in,
+                    layer_dim,
+                    downsample = is_first,
+                    expansion_rate = mbconv_expansion_rate,
+                    shrinkage_rate = mbconv_shrinkage_rate
+                )
+
+                block_attn = Attention(dim = layer_dim, dim_head = dim_head, dropout = dropout, window_size = window_size, num_registers = num_register_tokens)
+                block_ff = FeedForward(dim = layer_dim, dropout = dropout)
+
+                grid_attn = Attention(dim = layer_dim, dim_head = dim_head, dropout = dropout, window_size = window_size, num_registers = num_register_tokens)
+                grid_ff = FeedForward(dim = layer_dim, dropout = dropout)
+
+                register_tokens = nn.Parameter(torch.randn(num_register_tokens, layer_dim))
+
+                self.layers.append(ModuleList([
+                    conv,
+                    ModuleList([block_attn, block_ff]),
+                    ModuleList([grid_attn, grid_ff])
+                ]))
+
+                self.register_tokens.append(register_tokens)
+
+        # mlp head out
+
+        self.mlp_head = nn.Sequential(
+            Reduce('b d h w -> b d', 'mean'),
+            nn.LayerNorm(dims[-1]),
+            nn.Linear(dims[-1], num_classes)
+        )
+
+    def forward(self, x):
+        b, w = x.shape[0], self.window_size
+
+        x = self.conv_stem(x)
+
+        for (conv, (block_attn, block_ff), (grid_attn, grid_ff)), register_tokens in zip(self.layers, self.register_tokens):
+            x = conv(x)
+
+            # block-like attention
+
+            x = rearrange(x, 'b d (x w1) (y w2) -> b x y w1 w2 d', w1 = w, w2 = w)
+
+            # prepare register tokens
+
+            r = repeat(register_tokens, 'n d -> b x y n d', b = b, x = x.shape[1],y = x.shape[2])
+            r, register_batch_ps = pack_one(r, '* n d')
+
+            x, window_ps = pack_one(x, 'b x y * d')
+            x, batch_ps  = pack_one(x, '* n d')
+            x, register_ps = pack([r, x], 'b * d')
+
+            x = block_attn(x) + x
+            x = block_ff(x) + x
+
+            r, x = unpack(x, register_ps, 'b * d')
+
+            x = unpack_one(x, batch_ps, '* n d')
+            x = unpack_one(x, window_ps, 'b x y * d')
+            x = rearrange(x, 'b x y w1 w2 d -> b d (x w1) (y w2)')
+
+            r = unpack_one(r, register_batch_ps, '* n d')
+
+            # grid-like attention
+
+            x = rearrange(x, 'b d (w1 x) (w2 y) -> b x y w1 w2 d', w1 = w, w2 = w)
+
+            # prepare register tokens
+
+            r = reduce(r, 'b x y n d -> b n d', 'mean')
+            r = repeat(r, 'b n d -> b x y n d', x = x.shape[1], y = x.shape[2])
+            r, register_batch_ps = pack_one(r, '* n d')
+
+            x, window_ps = pack_one(x, 'b x y * d')
+            x, batch_ps  = pack_one(x, '* n d')
+            x, register_ps = pack([r, x], 'b * d')
+
+            x = grid_attn(x) + x
+
+            r, x = unpack(x, register_ps, 'b * d')
+
+            x = grid_ff(x) + x
+
+            x = unpack_one(x, batch_ps, '* n d')
+            x = unpack_one(x, window_ps, 'b x y * d')
+            x = rearrange(x, 'b x y w1 w2 d -> b d (w1 x) (w2 y)')
+
+        return self.mlp_head(x)

vit_pytorch/mobile_vit.py

@@ -1,252 +1,243 @@
-import torch
-import torch.nn as nn
-
-from einops import rearrange
-from einops.layers.torch import Reduce
-
-# helpers
-
-def conv_1x1_bn(inp, oup):
-    return nn.Sequential(
-        nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
-        nn.BatchNorm2d(oup),
-        nn.SiLU()
-    )
-
-def conv_nxn_bn(inp, oup, kernel_size=3, stride=1):
-    return nn.Sequential(
-        nn.Conv2d(inp, oup, kernel_size, stride, 1, bias=False),
-        nn.BatchNorm2d(oup),
-        nn.SiLU()
-    )
-
-# classes
-
-class PreNorm(nn.Module):
-    def __init__(self, dim, fn):
-        super().__init__()
-        self.norm = nn.LayerNorm(dim)
-        self.fn = fn
-
-    def forward(self, x, **kwargs):
-        return self.fn(self.norm(x), **kwargs)
-
-
-class FeedForward(nn.Module):
-    def __init__(self, dim, hidden_dim, dropout=0.):
-        super().__init__()
-        self.net = nn.Sequential(
-            nn.Linear(dim, hidden_dim),
-            nn.SiLU(),
-            nn.Dropout(dropout),
-            nn.Linear(hidden_dim, dim),
-            nn.Dropout(dropout)
-        )
-
-    def forward(self, x):
-        return self.net(x)
-
-class Attention(nn.Module):
-    def __init__(self, dim, heads=8, dim_head=64, dropout=0.):
-        super().__init__()
-        inner_dim = dim_head * heads
-        self.heads = heads
-        self.scale = dim_head ** -0.5
-
-        self.attend = nn.Softmax(dim=-1)
-        self.dropout = nn.Dropout(dropout)
-
-        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)
-
-        self.to_out = nn.Sequential(
-            nn.Linear(inner_dim, dim),
-            nn.Dropout(dropout)
-        )
-
-    def forward(self, x):
-        qkv = self.to_qkv(x).chunk(3, dim=-1)
-        q, k, v = map(lambda t: rearrange(
-            t, 'b p n (h d) -> b p h n d', h=self.heads), qkv)
-
-        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
-
-        attn = self.attend(dots)
-        attn = self.dropout(attn)
-
-        out = torch.matmul(attn, v)
-        out = rearrange(out, 'b p h n d -> b p n (h d)')
-        return self.to_out(out)
-
-class Transformer(nn.Module):
-    """Transformer block described in ViT.
-    Paper: https://arxiv.org/abs/2010.11929
-    Based on: https://github.com/lucidrains/vit-pytorch
-    """
-
-    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout=0.):
-        super().__init__()
-        self.layers = nn.ModuleList([])
-        for _ in range(depth):
-            self.layers.append(nn.ModuleList([
-                PreNorm(dim, Attention(dim, heads, dim_head, dropout)),
-                PreNorm(dim, FeedForward(dim, mlp_dim, dropout))
-            ]))
-
-    def forward(self, x):
-        for attn, ff in self.layers:
-            x = attn(x) + x
-            x = ff(x) + x
-        return x
-
-class MV2Block(nn.Module):
-    """MV2 block described in MobileNetV2.
-    Paper: https://arxiv.org/pdf/1801.04381
-    Based on: https://github.com/tonylins/pytorch-mobilenet-v2
-    """
-
-    def __init__(self, inp, oup, stride=1, expansion=4):
-        super().__init__()
-        self.stride = stride
-        assert stride in [1, 2]
-
-        hidden_dim = int(inp * expansion)
-        self.use_res_connect = self.stride == 1 and inp == oup
-
-        if expansion == 1:
-            self.conv = nn.Sequential(
-                # dw
-                nn.Conv2d(hidden_dim, hidden_dim, 3, stride,
-                          1, groups=hidden_dim, bias=False),
-                nn.BatchNorm2d(hidden_dim),
-                nn.SiLU(),
-                # pw-linear
-                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
-                nn.BatchNorm2d(oup),
-            )
-        else:
-            self.conv = nn.Sequential(
-                # pw
-                nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False),
-                nn.BatchNorm2d(hidden_dim),
-                nn.SiLU(),
-                # dw
-                nn.Conv2d(hidden_dim, hidden_dim, 3, stride,
-                          1, groups=hidden_dim, bias=False),
-                nn.BatchNorm2d(hidden_dim),
-                nn.SiLU(),
-                # pw-linear
-                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
-                nn.BatchNorm2d(oup),
-            )
-
-    def forward(self, x):
-        out = self.conv(x)
-        if self.use_res_connect:
-            out = out + x
-        return out
-
-class MobileViTBlock(nn.Module):
-    def __init__(self, dim, depth, channel, kernel_size, patch_size, mlp_dim, dropout=0.):
-        super().__init__()
-        self.ph, self.pw = patch_size
-
-        self.conv1 = conv_nxn_bn(channel, channel, kernel_size)
-        self.conv2 = conv_1x1_bn(channel, dim)
-
-        self.transformer = Transformer(dim, depth, 4, 8, mlp_dim, dropout)
-
-        self.conv3 = conv_1x1_bn(dim, channel)
-        self.conv4 = conv_nxn_bn(2 * channel, channel, kernel_size)
-
-    def forward(self, x):
-        y = x.clone()
-
-        # Local representations
-        x = self.conv1(x)
-        x = self.conv2(x)
-
-        # Global representations
-        _, _, h, w = x.shape
-        x = rearrange(x, 'b d (h ph) (w pw) -> b (ph pw) (h w) d',
-                      ph=self.ph, pw=self.pw)
-        x = self.transformer(x)
-        x = rearrange(x, 'b (ph pw) (h w) d -> b d (h ph) (w pw)',
-                      h=h//self.ph, w=w//self.pw, ph=self.ph, pw=self.pw)
-
-        # Fusion
-        x = self.conv3(x)
-        x = torch.cat((x, y), 1)
-        x = self.conv4(x)
-        return x
-
-class MobileViT(nn.Module):
-    """MobileViT.
-    Paper: https://arxiv.org/abs/2110.02178
-    Based on: https://github.com/chinhsuanwu/mobilevit-pytorch
-    """
-
-    def __init__(
-        self,
-        image_size,
-        dims,
-        channels,
-        num_classes,
-        expansion=4,
-        kernel_size=3,
-        patch_size=(2, 2),
-        depths=(2, 4, 3)
-    ):
-        super().__init__()
-        assert len(dims) == 3, 'dims must be a tuple of 3'
-        assert len(depths) == 3, 'depths must be a tuple of 3'
-
-        ih, iw = image_size
-        ph, pw = patch_size
-        assert ih % ph == 0 and iw % pw == 0
-
-        init_dim, *_, last_dim = channels
-
-        self.conv1 = conv_nxn_bn(3, init_dim, stride=2)
-
-        self.stem = nn.ModuleList([])
-        self.stem.append(MV2Block(channels[0], channels[1], 1, expansion))
-        self.stem.append(MV2Block(channels[1], channels[2], 2, expansion))
-        self.stem.append(MV2Block(channels[2], channels[3], 1, expansion))
-        self.stem.append(MV2Block(channels[2], channels[3], 1, expansion))
-
-        self.trunk = nn.ModuleList([])
-        self.trunk.append(nn.ModuleList([
-            MV2Block(channels[3], channels[4], 2, expansion),
-            MobileViTBlock(dims[0], depths[0], channels[5],
-                           kernel_size, patch_size, int(dims[0] * 2))
-        ]))
-
-        self.trunk.append(nn.ModuleList([
-            MV2Block(channels[5], channels[6], 2, expansion),
-            MobileViTBlock(dims[1], depths[1], channels[7],
-                           kernel_size, patch_size, int(dims[1] * 4))
-        ]))
-
-        self.trunk.append(nn.ModuleList([
-            MV2Block(channels[7], channels[8], 2, expansion),
-            MobileViTBlock(dims[2], depths[2], channels[9],
-                           kernel_size, patch_size, int(dims[2] * 4))
-        ]))
-
-        self.to_logits = nn.Sequential(
-            conv_1x1_bn(channels[-2], last_dim),
-            Reduce('b c h w -> b c', 'mean'),
-            nn.Linear(channels[-1], num_classes, bias=False)
-        )
-
-    def forward(self, x):
-        x = self.conv1(x)
-
-        for conv in self.stem:
-            x = conv(x)
-
-        for conv, attn in self.trunk:
-            x = conv(x)
-            x = attn(x)
-
-        return self.to_logits(x)
+import torch
+import torch.nn as nn
+
+from einops import rearrange
+from einops.layers.torch import Reduce
+
+# helpers
+
+def conv_1x1_bn(inp, oup):
+    return nn.Sequential(
+        nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
+        nn.BatchNorm2d(oup),
+        nn.SiLU()
+    )
+
+def conv_nxn_bn(inp, oup, kernel_size=3, stride=1):
+    return nn.Sequential(
+        nn.Conv2d(inp, oup, kernel_size, stride, 1, bias=False),
+        nn.BatchNorm2d(oup),
+        nn.SiLU()
+    )
+
+# classes
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, hidden_dim, dropout=0.):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Linear(dim, hidden_dim),
+            nn.SiLU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, dim),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+class Attention(nn.Module):
+    def __init__(self, dim, heads=8, dim_head=64, dropout=0.):
+        super().__init__()
+        inner_dim = dim_head * heads
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+
+        self.norm = nn.LayerNorm(dim)
+        self.attend = nn.Softmax(dim=-1)
+        self.dropout = nn.Dropout(dropout)
+
+        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, dim),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x):
+        x = self.norm(x)
+        qkv = self.to_qkv(x).chunk(3, dim=-1)
+
+        q, k, v = map(lambda t: rearrange(t, 'b p n (h d) -> b p h n d', h=self.heads), qkv)
+
+        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+
+        attn = self.attend(dots)
+        attn = self.dropout(attn)
+
+        out = torch.matmul(attn, v)
+        out = rearrange(out, 'b p h n d -> b p n (h d)')
+        return self.to_out(out)
+
+class Transformer(nn.Module):
+    """Transformer block described in ViT.
+    Paper: https://arxiv.org/abs/2010.11929
+    Based on: https://github.com/lucidrains/vit-pytorch
+    """
+
+    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout=0.):
+        super().__init__()
+        self.layers = nn.ModuleList([])
+        for _ in range(depth):
+            self.layers.append(nn.ModuleList([
+                Attention(dim, heads, dim_head, dropout),
+                FeedForward(dim, mlp_dim, dropout)
+            ]))
+
+    def forward(self, x):
+        for attn, ff in self.layers:
+            x = attn(x) + x
+            x = ff(x) + x
+        return x
+
+class MV2Block(nn.Module):
+    """MV2 block described in MobileNetV2.
+    Paper: https://arxiv.org/pdf/1801.04381
+    Based on: https://github.com/tonylins/pytorch-mobilenet-v2
+    """
+
+    def __init__(self, inp, oup, stride=1, expansion=4):
+        super().__init__()
+        self.stride = stride
+        assert stride in [1, 2]
+
+        hidden_dim = int(inp * expansion)
+        self.use_res_connect = self.stride == 1 and inp == oup
+
+        if expansion == 1:
+            self.conv = nn.Sequential(
+                # dw
+                nn.Conv2d(hidden_dim, hidden_dim, 3, stride,
+                          1, groups=hidden_dim, bias=False),
+                nn.BatchNorm2d(hidden_dim),
+                nn.SiLU(),
+                # pw-linear
+                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(oup),
+            )
+        else:
+            self.conv = nn.Sequential(
+                # pw
+                nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(hidden_dim),
+                nn.SiLU(),
+                # dw
+                nn.Conv2d(hidden_dim, hidden_dim, 3, stride,
+                          1, groups=hidden_dim, bias=False),
+                nn.BatchNorm2d(hidden_dim),
+                nn.SiLU(),
+                # pw-linear
+                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(oup),
+            )
+
+    def forward(self, x):
+        out = self.conv(x)
+        if self.use_res_connect:
+            out = out + x
+        return out
+
+class MobileViTBlock(nn.Module):
+    def __init__(self, dim, depth, channel, kernel_size, patch_size, mlp_dim, dropout=0.):
+        super().__init__()
+        self.ph, self.pw = patch_size
+
+        self.conv1 = conv_nxn_bn(channel, channel, kernel_size)
+        self.conv2 = conv_1x1_bn(channel, dim)
+
+        self.transformer = Transformer(dim, depth, 4, 8, mlp_dim, dropout)
+
+        self.conv3 = conv_1x1_bn(dim, channel)
+        self.conv4 = conv_nxn_bn(2 * channel, channel, kernel_size)
+
+    def forward(self, x):
+        y = x.clone()
+
+        # Local representations
+        x = self.conv1(x)
+        x = self.conv2(x)
+
+        # Global representations
+        _, _, h, w = x.shape
+        x = rearrange(x, 'b d (h ph) (w pw) -> b (ph pw) (h w) d', ph=self.ph, pw=self.pw)
+        x = self.transformer(x)
+        x = rearrange(x, 'b (ph pw) (h w) d -> b d (h ph) (w pw)', h=h//self.ph, w=w//self.pw, ph=self.ph, pw=self.pw)
+
+        # Fusion
+        x = self.conv3(x)
+        x = torch.cat((x, y), 1)
+        x = self.conv4(x)
+        return x
+
+class MobileViT(nn.Module):
+    """MobileViT.
+    Paper: https://arxiv.org/abs/2110.02178
+    Based on: https://github.com/chinhsuanwu/mobilevit-pytorch
+    """
+
+    def __init__(
+        self,
+        image_size,
+        dims,
+        channels,
+        num_classes,
+        expansion=4,
+        kernel_size=3,
+        patch_size=(2, 2),
+        depths=(2, 4, 3)
+    ):
+        super().__init__()
+        assert len(dims) == 3, 'dims must be a tuple of 3'
+        assert len(depths) == 3, 'depths must be a tuple of 3'
+
+        ih, iw = image_size
+        ph, pw = patch_size
+        assert ih % ph == 0 and iw % pw == 0
+
+        init_dim, *_, last_dim = channels
+
+        self.conv1 = conv_nxn_bn(3, init_dim, stride=2)
+
+        self.stem = nn.ModuleList([])
+        self.stem.append(MV2Block(channels[0], channels[1], 1, expansion))
+        self.stem.append(MV2Block(channels[1], channels[2], 2, expansion))
+        self.stem.append(MV2Block(channels[2], channels[3], 1, expansion))
+        self.stem.append(MV2Block(channels[2], channels[3], 1, expansion))
+
+        self.trunk = nn.ModuleList([])
+        self.trunk.append(nn.ModuleList([
+            MV2Block(channels[3], channels[4], 2, expansion),
+            MobileViTBlock(dims[0], depths[0], channels[5],
+                           kernel_size, patch_size, int(dims[0] * 2))
+        ]))
+
+        self.trunk.append(nn.ModuleList([
+            MV2Block(channels[5], channels[6], 2, expansion),
+            MobileViTBlock(dims[1], depths[1], channels[7],
+                           kernel_size, patch_size, int(dims[1] * 4))
+        ]))
+
+        self.trunk.append(nn.ModuleList([
+            MV2Block(channels[7], channels[8], 2, expansion),
+            MobileViTBlock(dims[2], depths[2], channels[9],
+                           kernel_size, patch_size, int(dims[2] * 4))
+        ]))
+
+        self.to_logits = nn.Sequential(
+            conv_1x1_bn(channels[-2], last_dim),
+            Reduce('b c h w -> b c', 'mean'),
+            nn.Linear(channels[-1], num_classes, bias=False)
+        )
+
+    def forward(self, x):
+        x = self.conv1(x)
+
+        for conv in self.stem:
+            x = conv(x)
+
+        for conv, attn in self.trunk:
+            x = conv(x)
+            x = attn(x)
+
+        return self.to_logits(x)

vit_pytorch/mp3.py

@@ -107,10 +107,10 @@ class ViT(nn.Module):
     def __init__(self, *, num_classes, image_size, patch_size, dim, depth, heads, mlp_dim, channels = 3, dim_head = 64, dropout = 0.):
         super().__init__()
         image_height, image_width = pair(image_size)
-        patch_height, patch_width = pair(patch_size)
+        self.patch_size = patch_height, patch_width = pair(patch_size)
 
         assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
-        
+
         num_patches = (image_height // patch_height) * (image_width // patch_width)
         patch_dim = channels * patch_height * patch_width
 
@@ -178,7 +178,7 @@ class MP3(nn.Module):
 
         attended_tokens = self.vit.transformer(tokens, tokens_unmasked)
         logits = rearrange(self.mlp_head(attended_tokens), 'b n d -> (b n) d')
-        
+
         # Define labels
         labels = repeat(torch.arange(num_patches, device = device), 'n -> (b n)', b = batch)
         loss = F.cross_entropy(logits, labels)

vit_pytorch/na_vit.py

@@ -1,4 +1,6 @@
-from functools import partial
+from __future__ import annotations
+
+from functools import partial, lru_cache
 from typing import List
 
 import torch
@@ -7,7 +9,6 @@ from torch import nn, Tensor
 from torch.nn.utils.rnn import pad_sequence as orig_pad_sequence
 
 from einops import rearrange, repeat
-from einops.layers.torch import Rearrange
 
 # helpers
 
@@ -17,12 +18,64 @@ def exists(val):
 def default(val, d):
     return val if exists(val) else d
 
+def always(val):
+    return lambda *args: val
+
 def pair(t):
     return t if isinstance(t, tuple) else (t, t)
 
 def divisible_by(numer, denom):
     return (numer % denom) == 0
 
+@lru_cache(maxsize=128)
+def posemb_grid(ph, pw, device):
+    h_idx = torch.arange(ph, device=device).repeat_interleave(pw)
+    w_idx = torch.arange(pw, device=device).repeat(ph)
+    return torch.stack([h_idx, w_idx], dim=-1)
+
+# auto grouping images
+
+def group_images_by_max_seq_len(
+    images: List[Tensor],
+    patch_size: int,
+    calc_token_dropout = None,
+    max_seq_len = 2048
+
+) -> List[List[Tensor]]:
+
+    calc_token_dropout = default(calc_token_dropout, always(0.))
+
+    groups = []
+    group = []
+    seq_len = 0
+
+    if isinstance(calc_token_dropout, (float, int)):
+        calc_token_dropout = always(calc_token_dropout)
+
+    for image in images:
+        assert isinstance(image, Tensor)
+
+        image_dims = image.shape[-2:]
+        ph, pw = map(lambda t: t // patch_size, image_dims)
+
+        image_seq_len = (ph * pw)
+        image_seq_len = int(image_seq_len * (1 - calc_token_dropout(*image_dims)))
+
+        assert image_seq_len <= max_seq_len, f'image with dimensions {image_dims} exceeds maximum sequence length'
+
+        if (seq_len + image_seq_len) > max_seq_len:
+            groups.append(group)
+            group = []
+            seq_len = 0
+
+        group.append(image)
+        seq_len += image_seq_len
+
+    if len(group) > 0:
+        groups.append(group)
+
+    return groups
+
 # normalization
 # they use layernorm without bias, something that pytorch does not offer
 
@@ -69,8 +122,7 @@ class Attention(nn.Module):
         self.q_norm = RMSNorm(heads, dim_head)
         self.k_norm = RMSNorm(heads, dim_head)
 
-        self.attend = nn.Softmax(dim = -1)
-        self.dropout = nn.Dropout(dropout)
+        self.dropout_p = dropout
 
         self.to_q = nn.Linear(dim, inner_dim, bias = False)
         self.to_kv = nn.Linear(dim, inner_dim * 2, bias = False)
@@ -97,19 +149,22 @@ class Attention(nn.Module):
         q = self.q_norm(q)
         k = self.k_norm(k)
 
-        dots = torch.matmul(q, k.transpose(-1, -2))
+        # combine masks if both exist
+        if exists(mask) or exists(attn_mask):
+            if exists(mask):
+                mask = rearrange(mask, 'b j -> b 1 1 j')
+            if exists(mask) and exists(attn_mask):
+                attn_mask = mask & attn_mask
+            elif exists(mask):
+                attn_mask = mask
 
-        if exists(mask):
-            mask = rearrange(mask, 'b j -> b 1 1 j')
-            dots = dots.masked_fill(~mask, -torch.finfo(dots.dtype).max)
+        out = F.scaled_dot_product_attention(
+            q, k, v,
+            attn_mask = attn_mask,
+            dropout_p = self.dropout_p if self.training else 0.,
+            scale = 1.  # RMSNorm already includes sqrt(dim) scaling
+        )
 
-        if exists(attn_mask):
-            dots = dots.masked_fill(~attn_mask, -torch.finfo(dots.dtype).max)
-
-        attn = self.attend(dots)
-        attn = self.dropout(attn)
-
-        out = torch.matmul(attn, v)
         out = rearrange(out, 'b h n d -> b n (h d)')
         return self.to_out(out)
 
@@ -146,13 +201,13 @@ class NaViT(nn.Module):
         # if int or float given, then assume constant dropout prob
         # otherwise accept a callback that in turn calculates dropout prob from height and width
 
-        self.calc_token_dropout = calc_token_dropout = None
+        self.calc_token_dropout = None
 
         if callable(token_dropout_prob):
             self.calc_token_dropout = token_dropout_prob
 
         elif isinstance(token_dropout_prob, (float, int)):
-            assert 0. < token_dropout_prob < 1.
+            assert 0. <= token_dropout_prob < 1.
             token_dropout_prob = float(token_dropout_prob)
             self.calc_token_dropout = lambda height, width: token_dropout_prob
 
@@ -199,13 +254,30 @@ class NaViT(nn.Module):
 
     def forward(
         self,
-        batched_images: List[List[Tensor]] # assume different resolution images already grouped correctly
+        batched_images: List[Tensor] | List[List[Tensor]], # assume different resolution images already grouped correctly
+        group_images = False,
+        group_max_seq_len = 2048
     ):
-        p, c, device, has_token_dropout = self.patch_size, self.channels, self.device, exists(self.calc_token_dropout)
+        p, c, device, has_token_dropout = self.patch_size, self.channels, self.device, exists(self.calc_token_dropout) and self.training
 
         arange = partial(torch.arange, device = device)
         pad_sequence = partial(orig_pad_sequence, batch_first = True)
 
+        # auto pack if specified
+
+        if group_images:
+            batched_images = group_images_by_max_seq_len(
+                batched_images,
+                patch_size = self.patch_size,
+                calc_token_dropout = self.calc_token_dropout if self.training else None,
+                max_seq_len = group_max_seq_len
+            )
+
+        # if List[Tensor] is not grouped -> List[List[Tensor]]
+
+        if torch.is_tensor(batched_images[0]):
+            batched_images = [batched_images]
+
         # process images into variable lengthed sequences with attention mask
 
         num_images = []
@@ -216,48 +288,47 @@ class NaViT(nn.Module):
         for images in batched_images:
             num_images.append(len(images))
 
-            sequences = []
-            positions = []
-            image_ids = torch.empty((0,), device = device, dtype = torch.long)
-
-            for image_id, image in enumerate(images):
-                assert image.ndim ==3 and image.shape[0] == c
+            # compute patch dimensions for all images
+            patch_dims = []
+            for image in images:
+                assert image.ndim == 3 and image.shape[0] == c
                 image_dims = image.shape[-2:]
                 assert all([divisible_by(dim, p) for dim in image_dims]), f'height and width {image_dims} of images must be divisible by patch size {p}'
+                patch_dims.append((image_dims[0] // p, image_dims[1] // p))
 
-                ph, pw = map(lambda dim: dim // p, image_dims)
+            # extract patches for all images
+            sequences = [rearrange(img, 'c (h p1) (w p2) -> (h w) (c p1 p2)', p1=p, p2=p) for img in images]
 
-                pos = torch.stack(torch.meshgrid((
-                    arange(ph),
-                    arange(pw)
-                ), indexing = 'ij'), dim = -1)
+            # compute positions - uses lru_cache to avoid redundant computation across forward passes
+            positions = [posemb_grid(ph, pw, device) for ph, pw in patch_dims]
 
-                pos = rearrange(pos, 'h w c -> (h w) c')
-                seq = rearrange(image, 'c (h p1) (w p2) -> (h w) (c p1 p2)', p1 = p, p2 = p)
-
-                seq_len = seq.shape[-2]
-
-                if has_token_dropout:
+            # handle token dropout
+            if has_token_dropout:
+                for i, (seq, pos) in enumerate(zip(sequences, positions)):
+                    image_dims = images[i].shape[-2:]
                     token_dropout = self.calc_token_dropout(*image_dims)
+                    seq_len = seq.shape[0]
                     num_keep = max(1, int(seq_len * (1 - token_dropout)))
-                    keep_indices = torch.randn((seq_len,), device = device).topk(num_keep, dim = -1).indices
+                    keep_indices = torch.randn((seq_len,), device=device).topk(num_keep, dim=-1).indices
+                    sequences[i] = seq[keep_indices]
+                    positions[i] = pos[keep_indices]
 
-                    seq = seq[keep_indices]
-                    pos = pos[keep_indices]
-
-                image_ids = F.pad(image_ids, (0, seq.shape[-2]), value = image_id)
-                sequences.append(seq)
-                positions.append(pos)
+            # build image_ids efficiently using repeat_interleave
+            patch_counts = [seq.shape[0] for seq in sequences]
+            image_ids = torch.repeat_interleave(
+                arange(len(images)),
+                torch.tensor(patch_counts, device=device)
+            )
 
             batched_image_ids.append(image_ids)
-            batched_sequences.append(torch.cat(sequences, dim = 0))
-            batched_positions.append(torch.cat(positions, dim = 0))
+            batched_sequences.append(torch.cat(sequences, dim=0))
+            batched_positions.append(torch.cat(positions, dim=0))
 
         # derive key padding mask
 
         lengths = torch.tensor([seq.shape[-2] for seq in batched_sequences], device = device, dtype = torch.long)
-        max_length = arange(lengths.amax().item())
-        key_pad_mask = rearrange(lengths, 'b -> b 1') <= rearrange(max_length, 'n -> 1 n')
+        seq_arange = arange(lengths.amax().item())
+        key_pad_mask = rearrange(seq_arange, 'n -> 1 n') < rearrange(lengths, 'b -> b 1')
 
         # derive attention mask, and combine with key padding mask from above
 
@@ -272,11 +343,11 @@ class NaViT(nn.Module):
 
         # need to know how many images for final attention pooling
 
-        num_images = torch.tensor(num_images, device = device, dtype = torch.long)        
+        num_images = torch.tensor(num_images, device = device, dtype = torch.long)
 
         # to patches
 
-        x = self.to_patch_embedding(patches)        
+        x = self.to_patch_embedding(patches)
 
         # factorized 2d absolute positional embedding
 

vit_pytorch/na_vit_nested_tensor.py

@@ -0,0 +1,330 @@
+from __future__ import annotations
+
+from typing import List
+from functools import partial
+
+import torch
+import packaging.version as pkg_version
+
+from torch import nn, Tensor
+import torch.nn.functional as F
+from torch.nn import Module, ModuleList
+from torch.nested import nested_tensor
+
+from einops import rearrange
+from einops.layers.torch import Rearrange
+
+# helpers
+
+def exists(val):
+    return val is not None
+
+def default(val, d):
+    return val if exists(val) else d
+
+def pair(t):
+    return t if isinstance(t, tuple) else (t, t)
+
+def divisible_by(numer, denom):
+    return (numer % denom) == 0
+
+# feedforward
+
+def FeedForward(dim, hidden_dim, dropout = 0.):
+    return nn.Sequential(
+        nn.LayerNorm(dim, bias = False),
+        nn.Linear(dim, hidden_dim),
+        nn.GELU(),
+        nn.Dropout(dropout),
+        nn.Linear(hidden_dim, dim),
+        nn.Dropout(dropout)
+    )
+
+class Attention(Module):
+    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0., qk_norm = True):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim, bias = False)
+
+        dim_inner = heads * dim_head
+        self.heads = heads
+        self.dim_head = dim_head
+
+        self.to_queries = nn.Linear(dim, dim_inner, bias = False)
+        self.to_keys = nn.Linear(dim, dim_inner, bias = False)
+        self.to_values = nn.Linear(dim, dim_inner, bias = False)
+
+        # in the paper, they employ qk rmsnorm, a way to stabilize attention
+        # will use layernorm in place of rmsnorm, which has been shown to work in certain papers. requires l2norm on non-ragged dimension to be supported in nested tensors
+
+        self.query_norm = nn.LayerNorm(dim_head, bias = False) if qk_norm else nn.Identity()
+        self.key_norm = nn.LayerNorm(dim_head, bias = False) if qk_norm else nn.Identity()
+
+        self.dropout = dropout
+
+        self.to_out = nn.Linear(dim_inner, dim, bias = False)
+
+    def forward(
+        self,
+        x,
+        context: Tensor | None = None
+    ):
+        x = self.norm(x)
+
+        # for attention pooling, one query pooling to entire sequence
+
+        context = default(context, x)
+
+        # queries, keys, values
+
+        query = self.to_queries(x)
+        key = self.to_keys(context)
+        value = self.to_values(context)
+
+        # split heads
+
+        def split_heads(t):
+            return t.unflatten(-1, (self.heads, self.dim_head))
+
+        def transpose_head_seq(t):
+            return t.transpose(1, 2)
+
+        query, key, value = map(split_heads, (query, key, value))
+
+        # qk norm for attention stability
+
+        query = self.query_norm(query)
+        key = self.key_norm(key)
+
+        query, key, value = map(transpose_head_seq, (query, key, value))
+
+        # attention
+
+        out = F.scaled_dot_product_attention(
+            query, key, value,
+            dropout_p = self.dropout if self.training else 0.
+        )
+
+        # merge heads
+
+        out = out.transpose(1, 2).flatten(-2)
+
+        return self.to_out(out)
+
+class Transformer(Module):
+    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0., qk_norm = True):
+        super().__init__()
+        self.layers = ModuleList([])
+
+        for _ in range(depth):
+            self.layers.append(ModuleList([
+                Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout, qk_norm = qk_norm),
+                FeedForward(dim, mlp_dim, dropout = dropout)
+            ]))
+
+        self.norm = nn.LayerNorm(dim, bias = False)
+
+    def forward(self, x):
+
+        for attn, ff in self.layers:
+            x = attn(x) + x
+            x = ff(x) + x
+
+        return self.norm(x)
+
+class NaViT(Module):
+    def __init__(
+        self,
+        *,
+        image_size,
+        patch_size,
+        num_classes,
+        dim,
+        depth,
+        heads,
+        mlp_dim,
+        channels = 3,
+        dim_head = 64,
+        dropout = 0.,
+        emb_dropout = 0.,
+        qk_rmsnorm = True,
+        token_dropout_prob: float | None = None
+    ):
+        super().__init__()
+
+        if pkg_version.parse(torch.__version__) < pkg_version.parse('2.5'):
+            print('nested tensor NaViT was tested on pytorch 2.5')
+
+
+        image_height, image_width = pair(image_size)
+
+        # what percent of tokens to dropout
+        # if int or float given, then assume constant dropout prob
+        # otherwise accept a callback that in turn calculates dropout prob from height and width
+
+        self.token_dropout_prob = token_dropout_prob
+
+        # calculate patching related stuff
+
+        assert divisible_by(image_height, patch_size) and divisible_by(image_width, patch_size), 'Image dimensions must be divisible by the patch size.'
+
+        patch_height_dim, patch_width_dim = (image_height // patch_size), (image_width // patch_size)
+        patch_dim = channels * (patch_size ** 2)
+
+        self.channels = channels
+        self.patch_size = patch_size
+        self.to_patches = Rearrange('c (h p1) (w p2) -> h w (c p1 p2)', p1 = patch_size, p2 = patch_size)
+
+        self.to_patch_embedding = nn.Sequential(
+            nn.LayerNorm(patch_dim),
+            nn.Linear(patch_dim, dim),
+            nn.LayerNorm(dim),
+        )
+
+        self.pos_embed_height = nn.Parameter(torch.randn(patch_height_dim, dim))
+        self.pos_embed_width = nn.Parameter(torch.randn(patch_width_dim, dim))
+
+        self.dropout = nn.Dropout(emb_dropout)
+
+        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout, qk_rmsnorm)
+
+        # final attention pooling queries
+
+        self.attn_pool_queries = nn.Parameter(torch.randn(dim))
+        self.attn_pool = Attention(dim = dim, dim_head = dim_head, heads = heads)
+
+        # output to logits
+
+        self.to_latent = nn.Identity()
+
+        self.mlp_head = nn.Sequential(
+            nn.LayerNorm(dim, bias = False),
+            nn.Linear(dim, num_classes, bias = False)
+        )
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    def forward(
+        self,
+        images: List[Tensor], # different resolution images
+    ):
+        batch, device = len(images), self.device
+        arange = partial(torch.arange, device = device)
+
+        assert all([image.ndim == 3 and image.shape[0] == self.channels for image in images]), f'all images must have {self.channels} channels and number of dimensions of 3 (channels, height, width)'
+
+        all_patches = [self.to_patches(image) for image in images]
+
+        # prepare factorized positional embedding height width indices
+
+        positions = []
+
+        for patches in all_patches:
+            patch_height, patch_width = patches.shape[:2]
+            hw_indices = torch.stack(torch.meshgrid((arange(patch_height), arange(patch_width)), indexing = 'ij'), dim = -1)
+            hw_indices = rearrange(hw_indices, 'h w c -> (h w) c')
+            positions.append(hw_indices)
+
+        # need the sizes to compute token dropout + positional embedding
+
+        tokens = [rearrange(patches, 'h w d -> (h w) d') for patches in all_patches]
+
+        # handle token dropout
+
+        seq_lens = torch.tensor([i.shape[0] for i in tokens], device = device)
+
+        if self.training and self.token_dropout_prob > 0:
+
+            keep_seq_lens = ((1. - self.token_dropout_prob) * seq_lens).int().clamp(min = 1)
+
+            kept_tokens = []
+            kept_positions = []
+
+            for one_image_tokens, one_image_positions, seq_len, num_keep in zip(tokens, positions, seq_lens, keep_seq_lens):
+                keep_indices = torch.randn((seq_len,), device = device).topk(num_keep, dim = -1).indices
+
+                one_image_kept_tokens = one_image_tokens[keep_indices]
+                one_image_kept_positions = one_image_positions[keep_indices]
+
+                kept_tokens.append(one_image_kept_tokens)
+                kept_positions.append(one_image_kept_positions)
+
+            tokens, positions, seq_lens = kept_tokens, kept_positions, keep_seq_lens
+
+        # add all height and width factorized positions
+
+        height_indices, width_indices = torch.cat(positions).unbind(dim = -1)
+        height_embed, width_embed = self.pos_embed_height[height_indices], self.pos_embed_width[width_indices]
+
+        pos_embed = height_embed + width_embed
+
+        # use nested tensor for transformers and save on padding computation
+
+        tokens = torch.cat(tokens)
+
+        # linear projection to patch embeddings
+
+        tokens = self.to_patch_embedding(tokens)
+
+        # absolute positions
+
+        tokens = tokens + pos_embed
+
+        tokens = nested_tensor(tokens.split(seq_lens.tolist()), layout = torch.jagged, device = device)
+
+        # embedding dropout
+
+        tokens = self.dropout(tokens)
+
+        # transformer
+
+        tokens = self.transformer(tokens)
+
+        # attention pooling
+        # will use a jagged tensor for queries, as SDPA requires all inputs to be jagged, or not
+
+        attn_pool_queries = [rearrange(self.attn_pool_queries, '... -> 1 ...')] * batch
+
+        attn_pool_queries = nested_tensor(attn_pool_queries, layout = torch.jagged)
+
+        pooled = self.attn_pool(attn_pool_queries, tokens)
+
+        # back to unjagged
+
+        logits = torch.stack(pooled.unbind())
+
+        logits = rearrange(logits, 'b 1 d -> b d')
+
+        logits = self.to_latent(logits)
+
+        return self.mlp_head(logits)
+
+# quick test
+
+if __name__ == '__main__':
+
+    v = NaViT(
+        image_size = 256,
+        patch_size = 32,
+        num_classes = 1000,
+        dim = 1024,
+        depth = 6,
+        heads = 16,
+        mlp_dim = 2048,
+        dropout = 0.,
+        emb_dropout = 0.,
+        token_dropout_prob = 0.1
+    )
+
+    # 5 images of different resolutions - List[Tensor]
+
+    images = [
+        torch.randn(3, 256, 256), torch.randn(3, 128, 128),
+        torch.randn(3, 128, 256), torch.randn(3, 256, 128),
+        torch.randn(3, 64, 256)
+    ]
+
+    assert v(images).shape == (5, 1000)
+
+    v(images).sum().backward()

vit_pytorch/na_vit_nested_tensor_3d.py

@@ -0,0 +1,356 @@
+from __future__ import annotations
+
+from typing import List
+from functools import partial
+
+import torch
+import packaging.version as pkg_version
+
+from torch import nn, Tensor
+import torch.nn.functional as F
+from torch.nn import Module, ModuleList
+from torch.nested import nested_tensor
+
+from einops import rearrange
+from einops.layers.torch import Rearrange
+
+# helpers
+
+def exists(val):
+    return val is not None
+
+def default(val, d):
+    return val if exists(val) else d
+
+def pair(t):
+    return t if isinstance(t, tuple) else (t, t)
+
+def divisible_by(numer, denom):
+    return (numer % denom) == 0
+
+# feedforward
+
+def FeedForward(dim, hidden_dim, dropout = 0.):
+    return nn.Sequential(
+        nn.LayerNorm(dim, bias = False),
+        nn.Linear(dim, hidden_dim),
+        nn.GELU(),
+        nn.Dropout(dropout),
+        nn.Linear(hidden_dim, dim),
+        nn.Dropout(dropout)
+    )
+
+class Attention(Module):
+    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0., qk_norm = True):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim, bias = False)
+
+        dim_inner = heads * dim_head
+        self.heads = heads
+        self.dim_head = dim_head
+
+        self.to_queries = nn.Linear(dim, dim_inner, bias = False)
+        self.to_keys = nn.Linear(dim, dim_inner, bias = False)
+        self.to_values = nn.Linear(dim, dim_inner, bias = False)
+
+        # in the paper, they employ qk rmsnorm, a way to stabilize attention
+        # will use layernorm in place of rmsnorm, which has been shown to work in certain papers. requires l2norm on non-ragged dimension to be supported in nested tensors
+
+        self.query_norm = nn.LayerNorm(dim_head, bias = False) if qk_norm else nn.Identity()
+        self.key_norm = nn.LayerNorm(dim_head, bias = False) if qk_norm else nn.Identity()
+
+        self.dropout = dropout
+
+        self.to_out = nn.Linear(dim_inner, dim, bias = False)
+
+    def forward(
+        self,
+        x,
+        context: Tensor | None = None
+    ):
+
+        x = self.norm(x)
+
+        # for attention pooling, one query pooling to entire sequence
+
+        context = default(context, x)
+
+        # queries, keys, values
+
+        query = self.to_queries(x)
+        key = self.to_keys(context)
+        value = self.to_values(context)
+
+        # split heads
+
+        def split_heads(t):
+            return t.unflatten(-1, (self.heads, self.dim_head))
+
+        def transpose_head_seq(t):
+            return t.transpose(1, 2)
+
+        query, key, value = map(split_heads, (query, key, value))
+
+        # qk norm for attention stability
+
+        query = self.query_norm(query)
+        key = self.key_norm(key)
+
+        query, key, value = map(transpose_head_seq, (query, key, value))
+
+        # attention
+
+        out = F.scaled_dot_product_attention(
+            query, key, value,
+            dropout_p = self.dropout if self.training else 0.
+        )
+
+        # merge heads
+
+        out = out.transpose(1, 2).flatten(-2)
+
+        return self.to_out(out)
+
+class Transformer(Module):
+    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0., qk_norm = True):
+        super().__init__()
+        self.layers = ModuleList([])
+
+        for _ in range(depth):
+            self.layers.append(ModuleList([
+                Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout, qk_norm = qk_norm),
+                FeedForward(dim, mlp_dim, dropout = dropout)
+            ]))
+
+        self.norm = nn.LayerNorm(dim, bias = False)
+
+    def forward(self, x):
+
+        for attn, ff in self.layers:
+            x = attn(x) + x
+            x = ff(x) + x
+
+        return self.norm(x)
+
+class NaViT(Module):
+    def __init__(
+        self,
+        *,
+        image_size,
+        max_frames,
+        patch_size,
+        frame_patch_size,
+        num_classes,
+        dim,
+        depth,
+        heads,
+        mlp_dim,
+        channels = 3,
+        dim_head = 64,
+        dropout = 0.,
+        emb_dropout = 0.,
+        num_registers = 4,
+        qk_rmsnorm = True,
+        token_dropout_prob: float | None = None
+    ):
+        super().__init__()
+        image_height, image_width = pair(image_size)
+
+        if pkg_version.parse(torch.__version__) < pkg_version.parse('2.5'):
+            print('nested tensor NaViT was tested on pytorch 2.5')
+
+        # what percent of tokens to dropout
+        # if int or float given, then assume constant dropout prob
+        # otherwise accept a callback that in turn calculates dropout prob from height and width
+
+        self.token_dropout_prob = token_dropout_prob
+
+        # calculate patching related stuff
+
+        assert divisible_by(image_height, patch_size) and divisible_by(image_width, patch_size), 'Image dimensions must be divisible by the patch size.'
+        assert divisible_by(max_frames, frame_patch_size)
+
+        patch_frame_dim, patch_height_dim, patch_width_dim = (max_frames // frame_patch_size), (image_height // patch_size), (image_width // patch_size)
+
+        patch_dim = channels * (patch_size ** 2) * frame_patch_size
+
+        self.channels = channels
+        self.patch_size = patch_size
+        self.to_patches = Rearrange('c (f pf) (h p1) (w p2) -> f h w (c pf p1 p2)', p1 = patch_size, p2 = patch_size, pf = frame_patch_size)
+
+        self.to_patch_embedding = nn.Sequential(
+            nn.LayerNorm(patch_dim),
+            nn.Linear(patch_dim, dim),
+            nn.LayerNorm(dim),
+        )
+
+        self.pos_embed_frame = nn.Parameter(torch.zeros(patch_frame_dim, dim))
+        self.pos_embed_height = nn.Parameter(torch.zeros(patch_height_dim, dim))
+        self.pos_embed_width = nn.Parameter(torch.zeros(patch_width_dim, dim))
+
+        # register tokens
+
+        self.register_tokens = nn.Parameter(torch.zeros(num_registers, dim))
+
+        nn.init.normal_(self.pos_embed_frame, std = 0.02)
+        nn.init.normal_(self.pos_embed_height, std = 0.02)
+        nn.init.normal_(self.pos_embed_width, std = 0.02)
+        nn.init.normal_(self.register_tokens, std = 0.02)
+
+        self.dropout = nn.Dropout(emb_dropout)
+
+        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout, qk_rmsnorm)
+
+        # final attention pooling queries
+
+        self.attn_pool_queries = nn.Parameter(torch.randn(dim))
+        self.attn_pool = Attention(dim = dim, dim_head = dim_head, heads = heads)
+
+        # output to logits
+
+        self.to_latent = nn.Identity()
+
+        self.mlp_head = nn.Sequential(
+            nn.LayerNorm(dim, bias = False),
+            nn.Linear(dim, num_classes, bias = False)
+        )
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    def forward(
+        self,
+        volumes: List[Tensor], # different resolution images / CT scans
+    ):
+        batch, device = len(volumes), self.device
+        arange = partial(torch.arange, device = device)
+
+        assert all([volume.ndim == 4 and volume.shape[0] == self.channels for volume in volumes]), f'all volumes must have {self.channels} channels and number of dimensions of {self.channels} (channels, frame, height, width)'
+
+        all_patches = [self.to_patches(volume) for volume in volumes]
+
+        # prepare factorized positional embedding height width indices
+
+        positions = []
+
+        for patches in all_patches:
+            patch_frame, patch_height, patch_width = patches.shape[:3]
+            fhw_indices = torch.stack(torch.meshgrid((arange(patch_frame), arange(patch_height), arange(patch_width)), indexing = 'ij'), dim = -1)
+            fhw_indices = rearrange(fhw_indices, 'f h w c -> (f h w) c')
+
+            positions.append(fhw_indices)
+
+        # need the sizes to compute token dropout + positional embedding
+
+        tokens = [rearrange(patches, 'f h w d -> (f h w) d') for patches in all_patches]
+
+        # handle token dropout
+
+        seq_lens = torch.tensor([i.shape[0] for i in tokens], device = device)
+
+        if self.training and self.token_dropout_prob > 0:
+
+            keep_seq_lens = ((1. - self.token_dropout_prob) * seq_lens).int().clamp(min = 1)
+
+            kept_tokens = []
+            kept_positions = []
+
+            for one_image_tokens, one_image_positions, seq_len, num_keep in zip(tokens, positions, seq_lens, keep_seq_lens):
+                keep_indices = torch.randn((seq_len,), device = device).topk(num_keep, dim = -1).indices
+
+                one_image_kept_tokens = one_image_tokens[keep_indices]
+                one_image_kept_positions = one_image_positions[keep_indices]
+
+                kept_tokens.append(one_image_kept_tokens)
+                kept_positions.append(one_image_kept_positions)
+
+            tokens, positions, seq_lens = kept_tokens, kept_positions, keep_seq_lens
+
+        # add all height and width factorized positions
+
+
+        frame_indices, height_indices, width_indices = torch.cat(positions).unbind(dim = -1)
+        frame_embed, height_embed, width_embed = self.pos_embed_frame[frame_indices], self.pos_embed_height[height_indices], self.pos_embed_width[width_indices]
+
+        pos_embed = frame_embed + height_embed + width_embed
+
+        tokens = torch.cat(tokens)
+
+        # linear projection to patch embeddings
+
+        tokens = self.to_patch_embedding(tokens)
+
+        # absolute positions
+
+        tokens = tokens + pos_embed
+
+        # add register tokens
+
+        tokens = tokens.split(seq_lens.tolist())
+
+        tokens = [torch.cat((self.register_tokens, one_tokens)) for one_tokens in tokens]
+
+        # use nested tensor for transformers and save on padding computation
+
+        tokens = nested_tensor(tokens, layout = torch.jagged, device = device)
+
+        # embedding dropout
+
+        tokens = self.dropout(tokens)
+
+        # transformer
+
+        tokens = self.transformer(tokens)
+
+        # attention pooling
+        # will use a jagged tensor for queries, as SDPA requires all inputs to be jagged, or not
+
+        attn_pool_queries = [rearrange(self.attn_pool_queries, '... -> 1 ...')] * batch
+
+        attn_pool_queries = nested_tensor(attn_pool_queries, layout = torch.jagged)
+
+        pooled = self.attn_pool(attn_pool_queries, tokens)
+
+        # back to unjagged
+
+        logits = torch.stack(pooled.unbind())
+
+        logits = rearrange(logits, 'b 1 d -> b d')
+
+        logits = self.to_latent(logits)
+
+        return self.mlp_head(logits)
+
+# quick test
+
+if __name__ == '__main__':
+
+    # works for torch 2.5
+
+    v = NaViT(
+        image_size = 256,
+        max_frames = 8,
+        patch_size = 32,
+        frame_patch_size = 2,
+        num_classes = 1000,
+        dim = 1024,
+        depth = 6,
+        heads = 16,
+        mlp_dim = 2048,
+        dropout = 0.,
+        emb_dropout = 0.,
+        token_dropout_prob = 0.1
+    )
+
+    # 5 volumetric data (videos or CT scans) of different resolutions - List[Tensor]
+
+    volumes = [
+        torch.randn(3, 2, 256, 256), torch.randn(3, 8, 128, 128),
+        torch.randn(3, 4, 128, 256), torch.randn(3, 2, 256, 128),
+        torch.randn(3, 4, 64, 256)
+    ]
+
+    assert v(volumes).shape == (5, 1000)
+
+    v(volumes).sum().backward()

vit_pytorch/nest.py

@@ -24,19 +24,11 @@ class LayerNorm(nn.Module):
         mean = torch.mean(x, dim = 1, keepdim = True)
         return (x - mean) / (var + self.eps).sqrt() * self.g + self.b
 
-class PreNorm(nn.Module):
-    def __init__(self, dim, fn):
-        super().__init__()
-        self.norm = LayerNorm(dim)
-        self.fn = fn
-
-    def forward(self, x, **kwargs):
-        return self.fn(self.norm(x), **kwargs)
-
 class FeedForward(nn.Module):
     def __init__(self, dim, mlp_mult = 4, dropout = 0.):
         super().__init__()
         self.net = nn.Sequential(
+            LayerNorm(dim),
             nn.Conv2d(dim, dim * mlp_mult, 1),
             nn.GELU(),
             nn.Dropout(dropout),
@@ -54,6 +46,7 @@ class Attention(nn.Module):
         self.heads = heads
         self.scale = dim_head ** -0.5
 
+        self.norm = LayerNorm(dim)
         self.attend = nn.Softmax(dim = -1)
         self.dropout = nn.Dropout(dropout)
         self.to_qkv = nn.Conv2d(dim, inner_dim * 3, 1, bias = False)
@@ -66,6 +59,8 @@ class Attention(nn.Module):
     def forward(self, x):
         b, c, h, w, heads = *x.shape, self.heads
 
+        x = self.norm(x)
+
         qkv = self.to_qkv(x).chunk(3, dim = 1)
         q, k, v = map(lambda t: rearrange(t, 'b (h d) x y -> b h (x y) d', h = heads), qkv)
 
@@ -93,8 +88,8 @@ class Transformer(nn.Module):
 
         for _ in range(depth):
             self.layers.append(nn.ModuleList([
-                PreNorm(dim, Attention(dim, heads = heads, dropout = dropout)),
-                PreNorm(dim, FeedForward(dim, mlp_mult, dropout = dropout))
+                Attention(dim, heads = heads, dropout = dropout),
+                FeedForward(dim, mlp_mult, dropout = dropout)
             ]))
     def forward(self, x):
         *_, h, w = x.shape

vit_pytorch/normalized_vit.py

@@ -0,0 +1,264 @@
+import torch
+from torch import nn
+from torch.nn import Module, ModuleList
+import torch.nn.functional as F
+import torch.nn.utils.parametrize as parametrize
+
+from einops import rearrange, reduce
+from einops.layers.torch import Rearrange
+
+# functions
+
+def exists(v):
+    return v is not None
+
+def default(v, d):
+    return v if exists(v) else d
+
+def pair(t):
+    return t if isinstance(t, tuple) else (t, t)
+
+def divisible_by(numer, denom):
+    return (numer % denom) == 0
+
+def l2norm(t, dim = -1):
+    return F.normalize(t, dim = dim, p = 2)
+
+# for use with parametrize
+
+class L2Norm(Module):
+    def __init__(self, dim = -1):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, t):
+        return l2norm(t, dim = self.dim)
+
+class NormLinear(Module):
+    def __init__(
+        self,
+        dim,
+        dim_out,
+        norm_dim_in = True
+    ):
+        super().__init__()
+        self.linear = nn.Linear(dim, dim_out, bias = False)
+
+        parametrize.register_parametrization(
+            self.linear,
+            'weight',
+            L2Norm(dim = -1 if norm_dim_in else 0)
+        )
+
+    @property
+    def weight(self):
+        return self.linear.weight
+
+    def forward(self, x):
+        return self.linear(x)
+
+# attention and feedforward
+
+class Attention(Module):
+    def __init__(
+        self,
+        dim,
+        *,
+        dim_head = 64,
+        heads = 8,
+        dropout = 0.
+    ):
+        super().__init__()
+        dim_inner = dim_head * heads
+        self.to_q = NormLinear(dim, dim_inner)
+        self.to_k = NormLinear(dim, dim_inner)
+        self.to_v = NormLinear(dim, dim_inner)
+
+        self.dropout = dropout
+
+        self.q_scale = nn.Parameter(torch.ones(heads, 1, dim_head) * (dim_head ** 0.25))
+        self.k_scale = nn.Parameter(torch.ones(heads, 1, dim_head) * (dim_head ** 0.25))
+
+        self.split_heads = Rearrange('b n (h d) -> b h n d', h = heads)
+        self.merge_heads = Rearrange('b h n d -> b n (h d)')
+
+        self.to_out = NormLinear(dim_inner, dim, norm_dim_in = False)
+
+    def forward(
+        self,
+        x
+    ):
+        q, k, v = self.to_q(x), self.to_k(x), self.to_v(x)
+
+        q, k, v = map(self.split_heads, (q, k, v))
+
+        # query key rmsnorm
+
+        q, k = map(l2norm, (q, k))
+
+        q = q * self.q_scale
+        k = k * self.k_scale
+
+        # scale is 1., as scaling factor is moved to s_qk (dk ^ 0.25) - eq. 16
+
+        out = F.scaled_dot_product_attention(
+            q, k, v,
+            dropout_p = self.dropout if self.training else 0.,
+            scale = 1.
+        )
+
+        out = self.merge_heads(out)
+        return self.to_out(out)
+
+class FeedForward(Module):
+    def __init__(
+        self,
+        dim,
+        *,
+        dim_inner,
+        dropout = 0.
+    ):
+        super().__init__()
+        dim_inner = int(dim_inner * 2 / 3)
+
+        self.dim = dim
+        self.dropout = nn.Dropout(dropout)
+
+        self.to_hidden = NormLinear(dim, dim_inner)
+        self.to_gate = NormLinear(dim, dim_inner)
+
+        self.hidden_scale = nn.Parameter(torch.ones(dim_inner))
+        self.gate_scale = nn.Parameter(torch.ones(dim_inner))
+
+        self.to_out = NormLinear(dim_inner, dim, norm_dim_in = False)
+
+    def forward(self, x):
+        hidden, gate = self.to_hidden(x), self.to_gate(x)
+
+        hidden = hidden * self.hidden_scale
+        gate = gate * self.gate_scale * (self.dim ** 0.5)
+
+        hidden = F.silu(gate) * hidden
+
+        hidden = self.dropout(hidden)
+        return self.to_out(hidden)
+
+# classes
+
+class nViT(Module):
+    """ https://arxiv.org/abs/2410.01131 """
+
+    def __init__(
+        self,
+        *,
+        image_size,
+        patch_size,
+        num_classes,
+        dim,
+        depth,
+        heads,
+        mlp_dim,
+        dropout = 0.,
+        channels = 3,
+        dim_head = 64,
+        residual_lerp_scale_init = None
+    ):
+        super().__init__()
+        image_height, image_width = pair(image_size)
+
+        # calculate patching related stuff
+
+        assert divisible_by(image_height, patch_size) and divisible_by(image_width, patch_size), 'Image dimensions must be divisible by the patch size.'
+
+        patch_height_dim, patch_width_dim = (image_height // patch_size), (image_width // patch_size)
+        patch_dim = channels * (patch_size ** 2)
+        num_patches = patch_height_dim * patch_width_dim
+
+        self.channels = channels
+        self.patch_size = patch_size
+
+        self.to_patch_embedding = nn.Sequential(
+            Rearrange('b c (h p1) (w p2) -> b (h w) (c p1 p2)', p1 = patch_size, p2 = patch_size),
+            NormLinear(patch_dim, dim, norm_dim_in = False),
+        )
+
+        self.abs_pos_emb = NormLinear(dim, num_patches)
+
+        residual_lerp_scale_init = default(residual_lerp_scale_init, 1. / depth)
+
+        # layers
+
+        self.dim = dim
+        self.scale = dim ** 0.5
+
+        self.layers = ModuleList([])
+        self.residual_lerp_scales = nn.ParameterList([])
+
+        for _ in range(depth):
+            self.layers.append(ModuleList([
+                Attention(dim, dim_head = dim_head, heads = heads, dropout = dropout),
+                FeedForward(dim, dim_inner = mlp_dim, dropout = dropout),
+            ]))
+
+            self.residual_lerp_scales.append(nn.ParameterList([
+                nn.Parameter(torch.ones(dim) * residual_lerp_scale_init / self.scale),
+                nn.Parameter(torch.ones(dim) * residual_lerp_scale_init / self.scale),
+            ]))
+
+        self.logit_scale = nn.Parameter(torch.ones(num_classes))
+
+        self.to_pred = NormLinear(dim, num_classes)
+
+    @torch.no_grad()
+    def norm_weights_(self):
+        for module in self.modules():
+            if not isinstance(module, NormLinear):
+                continue
+
+            normed = module.weight
+            original = module.linear.parametrizations.weight.original
+
+            original.copy_(normed)
+
+    def forward(self, images):
+        device = images.device
+
+        tokens = self.to_patch_embedding(images)
+
+        seq_len = tokens.shape[-2]
+        pos_emb = self.abs_pos_emb.weight[torch.arange(seq_len, device = device)]
+
+        tokens = l2norm(tokens + pos_emb)
+
+        for (attn, ff), (attn_alpha, ff_alpha) in zip(self.layers, self.residual_lerp_scales):
+
+            attn_out = l2norm(attn(tokens))
+            tokens = l2norm(tokens.lerp(attn_out, attn_alpha * self.scale))
+
+            ff_out = l2norm(ff(tokens))
+            tokens = l2norm(tokens.lerp(ff_out, ff_alpha * self.scale))
+
+        pooled = reduce(tokens, 'b n d -> b d', 'mean')
+
+        logits = self.to_pred(pooled)
+        logits = logits * self.logit_scale * self.scale
+
+        return logits
+
+# quick test
+
+if __name__ == '__main__':
+
+    v = nViT(
+        image_size = 256,
+        patch_size = 16,
+        num_classes = 1000,
+        dim = 1024,
+        depth = 6,
+        heads = 8,
+        mlp_dim = 2048,
+    )
+
+    img = torch.randn(4, 3, 256, 256)
+    logits = v(img) # (4, 1000)
+    assert logits.shape == (4, 1000)

vit_pytorch/parallel_vit.py

@@ -19,18 +19,11 @@ class Parallel(nn.Module):
     def forward(self, x):
         return sum([fn(x) for fn in self.fns])
 
-class PreNorm(nn.Module):
-    def __init__(self, dim, fn):
-        super().__init__()
-        self.norm = nn.LayerNorm(dim)
-        self.fn = fn
-    def forward(self, x, **kwargs):
-        return self.fn(self.norm(x), **kwargs)
-
 class FeedForward(nn.Module):
     def __init__(self, dim, hidden_dim, dropout = 0.):
         super().__init__()
         self.net = nn.Sequential(
+            nn.LayerNorm(dim),
             nn.Linear(dim, hidden_dim),
             nn.GELU(),
             nn.Dropout(dropout),
@@ -49,6 +42,7 @@ class Attention(nn.Module):
         self.heads = heads
         self.scale = dim_head ** -0.5
 
+        self.norm = nn.LayerNorm(dim)
         self.attend = nn.Softmax(dim = -1)
         self.dropout = nn.Dropout(dropout)
 
@@ -60,6 +54,7 @@ class Attention(nn.Module):
         ) if project_out else nn.Identity()
 
     def forward(self, x):
+        x = self.norm(x)
         qkv = self.to_qkv(x).chunk(3, dim = -1)
         q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
 
@@ -77,8 +72,8 @@ class Transformer(nn.Module):
         super().__init__()
         self.layers = nn.ModuleList([])
 
-        attn_block = lambda: PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout))
-        ff_block = lambda: PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout))        
+        attn_block = lambda: Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout)
+        ff_block = lambda: FeedForward(dim, mlp_dim, dropout = dropout)
 
         for _ in range(depth):
             self.layers.append(nn.ModuleList([
@@ -96,7 +91,7 @@ class ViT(nn.Module):
     def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, pool = 'cls', num_parallel_branches = 2, channels = 3, dim_head = 64, dropout = 0., emb_dropout = 0.):
         super().__init__()
         image_height, image_width = pair(image_size)
-        patch_height, patch_width = pair(patch_size)
+        self.patch_size = patch_height, patch_width = pair(patch_size)
 
         assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
 

vit_pytorch/pit.py

@@ -17,18 +17,11 @@ def conv_output_size(image_size, kernel_size, stride, padding = 0):
 
 # classes
 
-class PreNorm(nn.Module):
-    def __init__(self, dim, fn):
-        super().__init__()
-        self.norm = nn.LayerNorm(dim)
-        self.fn = fn
-    def forward(self, x, **kwargs):
-        return self.fn(self.norm(x), **kwargs)
-
 class FeedForward(nn.Module):
     def __init__(self, dim, hidden_dim, dropout = 0.):
         super().__init__()
         self.net = nn.Sequential(
+            nn.LayerNorm(dim),
             nn.Linear(dim, hidden_dim),
             nn.GELU(),
             nn.Dropout(dropout),
@@ -47,6 +40,7 @@ class Attention(nn.Module):
         self.heads = heads
         self.scale = dim_head ** -0.5
 
+        self.norm = nn.LayerNorm(dim)
         self.attend = nn.Softmax(dim = -1)
         self.dropout = nn.Dropout(dropout)
         self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
@@ -58,6 +52,8 @@ class Attention(nn.Module):
 
     def forward(self, x):
         b, n, _, h = *x.shape, self.heads
+
+        x = self.norm(x)
         qkv = self.to_qkv(x).chunk(3, dim = -1)
         q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), qkv)
 
@@ -76,8 +72,8 @@ class Transformer(nn.Module):
         self.layers = nn.ModuleList([])
         for _ in range(depth):
             self.layers.append(nn.ModuleList([
-                PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout)),
-                PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout))
+                Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout),
+                FeedForward(dim, mlp_dim, dropout = dropout)
             ]))
     def forward(self, x):
         for attn, ff in self.layers:
@@ -158,7 +154,7 @@ class PiT(nn.Module):
 
         for ind, (layer_depth, layer_heads) in enumerate(zip(depth, heads)):
             not_last = ind < (len(depth) - 1)
-            
+
             layers.append(Transformer(dim, layer_depth, layer_heads, dim_head, mlp_dim, dropout))
 
             if not_last:

vit_pytorch/regionvit.py

@@ -20,6 +20,18 @@ def divisible_by(val, d):
 
 # helper classes
 
+class ChanLayerNorm(nn.Module):
+    def __init__(self, dim, eps = 1e-5):
+        super().__init__()
+        self.eps = eps
+        self.g = nn.Parameter(torch.ones(1, dim, 1, 1))
+        self.b = nn.Parameter(torch.zeros(1, dim, 1, 1))
+
+    def forward(self, x):
+        var = torch.var(x, dim = 1, unbiased = False, keepdim = True)
+        mean = torch.mean(x, dim = 1, keepdim = True)
+        return (x - mean) / (var + self.eps).sqrt() * self.g + self.b
+
 class Downsample(nn.Module):
     def __init__(self, dim_in, dim_out):
         super().__init__()
@@ -134,7 +146,7 @@ class R2LTransformer(nn.Module):
         region_tokens = rearrange(region_tokens, 'b c h w -> b (h w) c')
 
         # calculate local relative positional bias
-        
+
         h_range = torch.arange(window_size_h, device = device)
         w_range = torch.arange(window_size_w, device = device)
 
@@ -212,10 +224,10 @@ class RegionViT(nn.Module):
         if tokenize_local_3_conv:
             self.local_encoder = nn.Sequential(
                 nn.Conv2d(3, init_dim, 3, 2, 1),
-                nn.LayerNorm(init_dim),
+                ChanLayerNorm(init_dim),
                 nn.GELU(),
                 nn.Conv2d(init_dim, init_dim, 3, 2, 1),
-                nn.LayerNorm(init_dim),
+                ChanLayerNorm(init_dim),
                 nn.GELU(),
                 nn.Conv2d(init_dim, init_dim, 3, 1, 1)
             )

vit_pytorch/rvt.py

@@ -3,12 +3,14 @@ from math import sqrt, pi, log
 import torch
 from torch import nn, einsum
 import torch.nn.functional as F
+from torch.amp import autocast
 
 from einops import rearrange, repeat
 from einops.layers.torch import Rearrange
 
 # rotary embeddings
 
+@autocast('cuda', enabled = False)
 def rotate_every_two(x):
     x = rearrange(x, '... (d j) -> ... d j', j = 2)
     x1, x2 = x.unbind(dim = -1)
@@ -22,6 +24,7 @@ class AxialRotaryEmbedding(nn.Module):
         scales = torch.linspace(1., max_freq / 2, self.dim // 4)
         self.register_buffer('scales', scales)
 
+    @autocast('cuda', enabled = False)
     def forward(self, x):
         device, dtype, n = x.device, x.dtype, int(sqrt(x.shape[-2]))
 
@@ -55,14 +58,6 @@ class DepthWiseConv2d(nn.Module):
 
 # helper classes
 
-class PreNorm(nn.Module):
-    def __init__(self, dim, fn):
-        super().__init__()
-        self.norm = nn.LayerNorm(dim)
-        self.fn = fn
-    def forward(self, x, **kwargs):
-        return self.fn(self.norm(x), **kwargs)
-
 class SpatialConv(nn.Module):
     def __init__(self, dim_in, dim_out, kernel, bias = False):
         super().__init__()
@@ -86,6 +81,7 @@ class FeedForward(nn.Module):
     def __init__(self, dim, hidden_dim, dropout = 0., use_glu = True):
         super().__init__()
         self.net = nn.Sequential(
+            nn.LayerNorm(dim),
             nn.Linear(dim, hidden_dim * 2 if use_glu else hidden_dim),
             GEGLU() if use_glu else nn.GELU(),
             nn.Dropout(dropout),
@@ -103,6 +99,7 @@ class Attention(nn.Module):
         self.heads = heads
         self.scale = dim_head ** -0.5
 
+        self.norm = nn.LayerNorm(dim)
         self.attend = nn.Softmax(dim = -1)
         self.dropout = nn.Dropout(dropout)
 
@@ -121,6 +118,9 @@ class Attention(nn.Module):
         b, n, _, h = *x.shape, self.heads
 
         to_q_kwargs = {'fmap_dims': fmap_dims} if self.use_ds_conv else {}
+
+        x = self.norm(x)
+
         q = self.to_q(x, **to_q_kwargs)
 
         qkv = (q, *self.to_kv(x).chunk(2, dim = -1))
@@ -162,8 +162,8 @@ class Transformer(nn.Module):
         self.pos_emb = AxialRotaryEmbedding(dim_head, max_freq = image_size)
         for _ in range(depth):
             self.layers.append(nn.ModuleList([
-                PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout, use_rotary = use_rotary, use_ds_conv = use_ds_conv)),
-                PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout, use_glu = use_glu))
+                Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout, use_rotary = use_rotary, use_ds_conv = use_ds_conv),
+                FeedForward(dim, mlp_dim, dropout = dropout, use_glu = use_glu)
             ]))
     def forward(self, x, fmap_dims):
         pos_emb = self.pos_emb(x[:, 1:])

vit_pytorch/scalable_vit.py

@@ -33,15 +33,6 @@ class ChanLayerNorm(nn.Module):
         mean = torch.mean(x, dim = 1, keepdim = True)
         return (x - mean) / (var + self.eps).sqrt() * self.g + self.b
 
-class PreNorm(nn.Module):
-    def __init__(self, dim, fn):
-        super().__init__()
-        self.norm = ChanLayerNorm(dim)
-        self.fn = fn
-
-    def forward(self, x):
-        return self.fn(self.norm(x))
-
 class Downsample(nn.Module):
     def __init__(self, dim_in, dim_out):
         super().__init__()
@@ -65,6 +56,7 @@ class FeedForward(nn.Module):
         super().__init__()
         inner_dim = dim * expansion_factor
         self.net = nn.Sequential(
+            ChanLayerNorm(dim),
             nn.Conv2d(dim, inner_dim, 1),
             nn.GELU(),
             nn.Dropout(dropout),
@@ -92,6 +84,7 @@ class ScalableSelfAttention(nn.Module):
         self.attend = nn.Softmax(dim = -1)
         self.dropout = nn.Dropout(dropout)
 
+        self.norm = ChanLayerNorm(dim)
         self.to_q = nn.Conv2d(dim, dim_key * heads, 1, bias = False)
         self.to_k = nn.Conv2d(dim, dim_key * heads, reduction_factor, stride = reduction_factor, bias = False)
         self.to_v = nn.Conv2d(dim, dim_value * heads, reduction_factor, stride = reduction_factor, bias = False)
@@ -104,6 +97,8 @@ class ScalableSelfAttention(nn.Module):
     def forward(self, x):
         height, width, heads = *x.shape[-2:], self.heads
 
+        x = self.norm(x)
+
         q, k, v = self.to_q(x), self.to_k(x), self.to_v(x)
 
         # split out heads
@@ -145,6 +140,7 @@ class InteractiveWindowedSelfAttention(nn.Module):
         self.attend = nn.Softmax(dim = -1)
         self.dropout = nn.Dropout(dropout)
 
+        self.norm = ChanLayerNorm(dim)
         self.local_interactive_module = nn.Conv2d(dim_value * heads, dim_value * heads, 3, padding = 1)
 
         self.to_q = nn.Conv2d(dim, dim_key * heads, 1, bias = False)
@@ -159,6 +155,8 @@ class InteractiveWindowedSelfAttention(nn.Module):
     def forward(self, x):
         height, width, heads, wsz = *x.shape[-2:], self.heads, self.window_size
 
+        x = self.norm(x)
+
         wsz_h, wsz_w = default(wsz, height), default(wsz, width)
         assert (height % wsz_h) == 0 and (width % wsz_w) == 0, f'height ({height}) or width ({width}) of feature map is not divisible by the window size ({wsz_h}, {wsz_w})'
 
@@ -189,7 +187,7 @@ class InteractiveWindowedSelfAttention(nn.Module):
 
         out = rearrange(out, '(b x y) h (w1 w2) d -> b (h d) (x w1) (y w2)', x = height // wsz_h, y = width // wsz_w, w1 = wsz_h, w2 = wsz_w)
 
-        # add LIM output 
+        # add LIM output
 
         out = out + local_out
 
@@ -217,11 +215,11 @@ class Transformer(nn.Module):
             is_first = ind == 0
 
             self.layers.append(nn.ModuleList([
-                PreNorm(dim, ScalableSelfAttention(dim, heads = heads, dim_key = ssa_dim_key, dim_value = ssa_dim_value, reduction_factor = ssa_reduction_factor, dropout = dropout)),
-                PreNorm(dim, FeedForward(dim, expansion_factor = ff_expansion_factor, dropout = dropout)),
+                ScalableSelfAttention(dim, heads = heads, dim_key = ssa_dim_key, dim_value = ssa_dim_value, reduction_factor = ssa_reduction_factor, dropout = dropout),
+                FeedForward(dim, expansion_factor = ff_expansion_factor, dropout = dropout),
                 PEG(dim) if is_first else None,
-                PreNorm(dim, FeedForward(dim, expansion_factor = ff_expansion_factor, dropout = dropout)),
-                PreNorm(dim, InteractiveWindowedSelfAttention(dim, heads = heads, dim_key = iwsa_dim_key, dim_value = iwsa_dim_value, window_size = iwsa_window_size, dropout = dropout))
+                FeedForward(dim, expansion_factor = ff_expansion_factor, dropout = dropout),
+                InteractiveWindowedSelfAttention(dim, heads = heads, dim_key = iwsa_dim_key, dim_value = iwsa_dim_value, window_size = iwsa_window_size, dropout = dropout)
             ]))
 
         self.norm = ChanLayerNorm(dim) if norm_output else nn.Identity()

vit_pytorch/sep_vit.py

@@ -25,15 +25,6 @@ class ChanLayerNorm(nn.Module):
         mean = torch.mean(x, dim = 1, keepdim = True)
         return (x - mean) / (var + self.eps).sqrt() * self.g + self.b
 
-class PreNorm(nn.Module):
-    def __init__(self, dim, fn):
-        super().__init__()
-        self.norm = ChanLayerNorm(dim)
-        self.fn = fn
-
-    def forward(self, x):
-        return self.fn(self.norm(x))
-
 class OverlappingPatchEmbed(nn.Module):
     def __init__(self, dim_in, dim_out, stride = 2):
         super().__init__()
@@ -59,6 +50,7 @@ class FeedForward(nn.Module):
         super().__init__()
         inner_dim = int(dim * mult)
         self.net = nn.Sequential(
+            ChanLayerNorm(dim),
             nn.Conv2d(dim, inner_dim, 1),
             nn.GELU(),
             nn.Dropout(dropout),
@@ -85,6 +77,8 @@ class DSSA(nn.Module):
         self.window_size = window_size
         inner_dim = dim_head * heads
 
+        self.norm = ChanLayerNorm(dim)
+
         self.attend = nn.Sequential(
             nn.Softmax(dim = -1),
             nn.Dropout(dropout)
@@ -138,6 +132,8 @@ class DSSA(nn.Module):
         assert (height % wsz) == 0 and (width % wsz) == 0, f'height {height} and width {width} must be divisible by window size {wsz}'
         num_windows = (height // wsz) * (width // wsz)
 
+        x = self.norm(x)
+
         # fold in windows for "depthwise" attention - not sure why it is named depthwise when it is just "windowed" attention
 
         x = rearrange(x, 'b c (h w1) (w w2) -> (b h w) c (w1 w2)', w1 = wsz, w2 = wsz)
@@ -225,8 +221,8 @@ class Transformer(nn.Module):
 
         for ind in range(depth):
             self.layers.append(nn.ModuleList([
-                PreNorm(dim, DSSA(dim, heads = heads, dim_head = dim_head, dropout = dropout)),
-                PreNorm(dim, FeedForward(dim, mult = ff_mult, dropout = dropout)),
+                DSSA(dim, heads = heads, dim_head = dim_head, dropout = dropout),
+                FeedForward(dim, mult = ff_mult, dropout = dropout),
             ]))
 
         self.norm = ChanLayerNorm(dim) if norm_output else nn.Identity()

vit_pytorch/simple_flash_attn_vit.py

@@ -26,7 +26,7 @@ def posemb_sincos_2d(patches, temperature = 10000, dtype = torch.float32):
     omega = 1. / (temperature ** omega)
 
     y = y.flatten()[:, None] * omega[None, :]
-    x = x.flatten()[:, None] * omega[None, :] 
+    x = x.flatten()[:, None] * omega[None, :]
     pe = torch.cat((x.sin(), x.cos(), y.sin(), y.cos()), dim = 1)
     return pe.type(dtype)
 
@@ -57,7 +57,7 @@ class Attend(nn.Module):
         config = self.cuda_config if q.is_cuda else self.cpu_config
 
         # flash attention - https://arxiv.org/abs/2205.14135
-        
+
         with torch.backends.cuda.sdp_kernel(**config._asdict()):
             out = F.scaled_dot_product_attention(q, k, v)
 
@@ -140,7 +140,7 @@ class SimpleViT(nn.Module):
     def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, channels = 3, dim_head = 64, use_flash = True):
         super().__init__()
         image_height, image_width = pair(image_size)
-        patch_height, patch_width = pair(patch_size)
+        self.patch_size = patch_height, patch_width = pair(patch_size)
 
         assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
 

vit_pytorch/simple_flash_attn_vit_3d.py

@@ -0,0 +1,171 @@
+from packaging import version
+from collections import namedtuple
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.nn import Module, ModuleList
+
+from einops import rearrange
+from einops.layers.torch import Rearrange
+
+# constants
+
+Config = namedtuple('FlashAttentionConfig', ['enable_flash', 'enable_math', 'enable_mem_efficient'])
+
+# helpers
+
+def pair(t):
+    return t if isinstance(t, tuple) else (t, t)
+
+def posemb_sincos_3d(patches, temperature = 10000, dtype = torch.float32):
+    _, f, h, w, dim, device, dtype = *patches.shape, patches.device, patches.dtype
+
+    z, y, x = torch.meshgrid(
+        torch.arange(f, device = device),
+        torch.arange(h, device = device),
+        torch.arange(w, device = device),
+    indexing = 'ij')
+
+    fourier_dim = dim // 6
+
+    omega = torch.arange(fourier_dim, device = device) / (fourier_dim - 1)
+    omega = 1. / (temperature ** omega)
+
+    z = z.flatten()[:, None] * omega[None, :]
+    y = y.flatten()[:, None] * omega[None, :]
+    x = x.flatten()[:, None] * omega[None, :]
+
+    pe = torch.cat((x.sin(), x.cos(), y.sin(), y.cos(), z.sin(), z.cos()), dim = 1)
+
+    pe = F.pad(pe, (0, dim - (fourier_dim * 6))) # pad if feature dimension not cleanly divisible by 6
+    return pe.type(dtype)
+
+# main class
+
+class Attend(Module):
+    def __init__(self, use_flash = False, config: Config = Config(True, True, True)):
+        super().__init__()
+        self.config = config
+        self.use_flash = use_flash
+        assert not (use_flash and version.parse(torch.__version__) < version.parse('2.0.0')), 'in order to use flash attention, you must be using pytorch 2.0 or above'
+
+    def flash_attn(self, q, k, v):
+        # flash attention - https://arxiv.org/abs/2205.14135
+
+        with torch.backends.cuda.sdp_kernel(**self.config._asdict()):
+            out = F.scaled_dot_product_attention(q, k, v)
+
+        return out
+
+    def forward(self, q, k, v):
+        n, device, scale = q.shape[-2], q.device, q.shape[-1] ** -0.5
+
+        if self.use_flash:
+            return self.flash_attn(q, k, v)
+
+        # similarity
+
+        sim = einsum("b h i d, b j d -> b h i j", q, k) * scale
+
+        # attention
+
+        attn = sim.softmax(dim=-1)
+
+        # aggregate values
+
+        out = einsum("b h i j, b j d -> b h i d", attn, v)
+
+        return out
+
+# classes
+
+class FeedForward(Module):
+    def __init__(self, dim, hidden_dim):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Linear(dim, hidden_dim),
+            nn.GELU(),
+            nn.Linear(hidden_dim, dim),
+        )
+    def forward(self, x):
+        return self.net(x)
+
+class Attention(Module):
+    def __init__(self, dim, heads = 8, dim_head = 64, use_flash = True):
+        super().__init__()
+        inner_dim = dim_head *  heads
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+        self.norm = nn.LayerNorm(dim)
+
+        self.attend = Attend(use_flash = use_flash)
+
+        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
+        self.to_out = nn.Linear(inner_dim, dim, bias = False)
+
+    def forward(self, x):
+        x = self.norm(x)
+
+        qkv = self.to_qkv(x).chunk(3, dim = -1)
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
+
+        out = self.attend(q, k, v)
+
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)
+
+class Transformer(Module):
+    def __init__(self, dim, depth, heads, dim_head, mlp_dim, use_flash):
+        super().__init__()
+        self.layers = ModuleList([])
+        for _ in range(depth):
+            self.layers.append(ModuleList([
+                Attention(dim, heads = heads, dim_head = dim_head, use_flash = use_flash),
+                FeedForward(dim, mlp_dim)
+            ]))
+
+    def forward(self, x):
+        for attn, ff in self.layers:
+            x = attn(x) + x
+            x = ff(x) + x
+
+        return x
+
+class SimpleViT(Module):
+    def __init__(self, *, image_size, image_patch_size, frames, frame_patch_size, num_classes, dim, depth, heads, mlp_dim, channels = 3, dim_head = 64, use_flash_attn = True):
+        super().__init__()
+        image_height, image_width = pair(image_size)
+        self.patch_size = patch_height, patch_width = pair(image_patch_size)
+
+        assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
+        assert frames % frame_patch_size == 0, 'Frames must be divisible by the frame patch size'
+
+        num_patches = (image_height // patch_height) * (image_width // patch_width) * (frames // frame_patch_size)
+        patch_dim = channels * patch_height * patch_width * frame_patch_size
+
+        self.to_patch_embedding = nn.Sequential(
+            Rearrange('b c (f pf) (h p1) (w p2) -> b f h w (pf p1 p2 c)', p1 = patch_height, p2 = patch_width, pf = frame_patch_size),
+            nn.LayerNorm(patch_dim),
+            nn.Linear(patch_dim, dim),
+            nn.LayerNorm(dim),
+        )
+
+        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, use_flash_attn)
+
+        self.to_latent = nn.Identity()
+        self.linear_head = nn.Linear(dim, num_classes)
+
+    def forward(self, video):
+        *_, h, w, dtype = *video.shape, video.dtype
+
+        x = self.to_patch_embedding(video)
+        pe = posemb_sincos_3d(x)
+        x = rearrange(x, 'b ... d -> b (...) d') + pe
+
+        x = self.transformer(x)
+        x = x.mean(dim = 1)
+
+        x = self.to_latent(x)
+        return self.linear_head(x)

vit_pytorch/simple_uvit.py

@@ -0,0 +1,176 @@
+import torch
+from torch import nn
+from torch.nn import Module, ModuleList
+
+from einops import rearrange, repeat, pack, unpack
+from einops.layers.torch import Rearrange
+
+# helpers
+
+def pair(t):
+    return t if isinstance(t, tuple) else (t, t)
+
+def exists(v):
+    return v is not None
+
+def divisible_by(num, den):
+    return (num % den) == 0
+
+def posemb_sincos_2d(h, w, dim, temperature: int = 10000, dtype = torch.float32):
+    y, x = torch.meshgrid(torch.arange(h), torch.arange(w), indexing="ij")
+    assert divisible_by(dim, 4), "feature dimension must be multiple of 4 for sincos emb"
+    omega = torch.arange(dim // 4) / (dim // 4 - 1)
+    omega = temperature ** -omega
+
+    y = y.flatten()[:, None] * omega[None, :]
+    x = x.flatten()[:, None] * omega[None, :]
+    pe = torch.cat((x.sin(), x.cos(), y.sin(), y.cos()), dim=1)
+    return pe.type(dtype)
+
+# classes
+
+def FeedForward(dim, hidden_dim):
+    return nn.Sequential(
+        nn.LayerNorm(dim),
+        nn.Linear(dim, hidden_dim),
+        nn.GELU(),
+        nn.Linear(hidden_dim, dim),
+    )
+
+class Attention(Module):
+    def __init__(self, dim, heads = 8, dim_head = 64):
+        super().__init__()
+        inner_dim = dim_head *  heads
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+        self.norm = nn.LayerNorm(dim)
+
+        self.attend = nn.Softmax(dim = -1)
+
+        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
+        self.to_out = nn.Linear(inner_dim, dim, bias = False)
+
+    def forward(self, x):
+        x = self.norm(x)
+
+        qkv = self.to_qkv(x).chunk(3, dim = -1)
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
+
+        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+
+        attn = self.attend(dots)
+
+        out = torch.matmul(attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)
+
+class Transformer(Module):
+    def __init__(self, dim, depth, heads, dim_head, mlp_dim):
+        super().__init__()
+        self.depth = depth
+        self.norm = nn.LayerNorm(dim)
+        self.layers = ModuleList([])
+
+        for layer in range(1, depth + 1):
+            latter_half = layer >= (depth / 2 + 1)
+
+            self.layers.append(nn.ModuleList([
+                nn.Linear(dim * 2, dim) if latter_half else None,
+                Attention(dim, heads = heads, dim_head = dim_head),
+                FeedForward(dim, mlp_dim)
+            ]))
+
+    def forward(self, x):
+
+        skips = []
+
+        for ind, (combine_skip, attn, ff) in enumerate(self.layers):
+            layer = ind + 1
+            first_half = layer <= (self.depth / 2)
+
+            if first_half:
+                skips.append(x)
+
+            if exists(combine_skip):
+                skip = skips.pop()
+                skip_and_x = torch.cat((skip, x), dim = -1)
+                x = combine_skip(skip_and_x)
+
+            x = attn(x) + x
+            x = ff(x) + x
+
+        assert len(skips) == 0
+
+        return self.norm(x)
+
+class SimpleUViT(Module):
+    def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, num_register_tokens = 4, channels = 3, dim_head = 64):
+        super().__init__()
+        image_height, image_width = pair(image_size)
+        self.patch_size = patch_height, patch_width = pair(patch_size)
+
+        assert divisible_by(image_height, patch_height) and divisible_by(image_width, patch_width), 'Image dimensions must be divisible by the patch size.'
+
+        patch_dim = channels * patch_height * patch_width
+
+        self.to_patch_embedding = nn.Sequential(
+            Rearrange("b c (h p1) (w p2) -> b (h w) (p1 p2 c)", p1 = patch_height, p2 = patch_width),
+            nn.LayerNorm(patch_dim),
+            nn.Linear(patch_dim, dim),
+            nn.LayerNorm(dim),
+        )
+
+        pos_embedding = posemb_sincos_2d(
+            h = image_height // patch_height,
+            w = image_width // patch_width,
+            dim = dim
+        )
+
+        self.register_buffer('pos_embedding', pos_embedding, persistent = False)
+
+        self.register_tokens = nn.Parameter(torch.randn(num_register_tokens, dim))
+
+        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim)
+
+        self.pool = "mean"
+        self.to_latent = nn.Identity()
+
+        self.linear_head = nn.Linear(dim, num_classes)
+
+    def forward(self, img):
+        batch, device = img.shape[0], img.device
+
+        x = self.to_patch_embedding(img)
+        x = x + self.pos_embedding.type(x.dtype)
+
+        r = repeat(self.register_tokens, 'n d -> b n d', b = batch)
+
+        x, ps = pack([x, r], 'b * d')
+
+        x = self.transformer(x)
+
+        x, _ = unpack(x, ps, 'b * d')
+
+        x = x.mean(dim = 1)
+
+        x = self.to_latent(x)
+        return self.linear_head(x)
+
+# quick test on odd number of layers
+
+if __name__ == '__main__':
+
+    v = SimpleUViT(
+        image_size = 256,
+        patch_size = 32,
+        num_classes = 1000,
+        dim = 1024,
+        depth = 7,
+        heads = 16,
+        mlp_dim = 2048
+    ).cuda()
+
+    img = torch.randn(2, 3, 256, 256).cuda()
+
+    preds = v(img)
+    assert preds.shape == (2, 1000)

vit_pytorch/simple_vit.py

@@ -64,6 +64,7 @@ class Attention(nn.Module):
 class Transformer(nn.Module):
     def __init__(self, dim, depth, heads, dim_head, mlp_dim):
         super().__init__()
+        self.norm = nn.LayerNorm(dim)
         self.layers = nn.ModuleList([])
         for _ in range(depth):
             self.layers.append(nn.ModuleList([
@@ -74,13 +75,13 @@ class Transformer(nn.Module):
         for attn, ff in self.layers:
             x = attn(x) + x
             x = ff(x) + x
-        return x
+        return self.norm(x)
 
 class SimpleViT(nn.Module):
     def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, channels = 3, dim_head = 64):
         super().__init__()
         image_height, image_width = pair(image_size)
-        patch_height, patch_width = pair(patch_size)
+        self.patch_size = patch_height, patch_width = pair(patch_size)
 
         assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
 
@@ -97,16 +98,14 @@ class SimpleViT(nn.Module):
             h = image_height // patch_height,
             w = image_width // patch_width,
             dim = dim,
-        ) 
+        )
 
         self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim)
 
-        self.to_latent = nn.Identity()
-        self.linear_head = nn.Sequential(
-            nn.LayerNorm(dim),
-            nn.Linear(dim, num_classes)
-        )
         self.pool = "mean"
+        self.to_latent = nn.Identity()
+
+        self.linear_head = nn.Linear(dim, num_classes)
 
     def forward(self, img):
         device = img.device

vit_pytorch/simple_vit_1d.py

@@ -62,6 +62,7 @@ class Attention(nn.Module):
 class Transformer(nn.Module):
     def __init__(self, dim, depth, heads, dim_head, mlp_dim):
         super().__init__()
+        self.norm = nn.LayerNorm(dim)
         self.layers = nn.ModuleList([])
         for _ in range(depth):
             self.layers.append(nn.ModuleList([
@@ -72,7 +73,7 @@ class Transformer(nn.Module):
         for attn, ff in self.layers:
             x = attn(x) + x
             x = ff(x) + x
-        return x
+        return self.norm(x)
 
 class SimpleViT(nn.Module):
     def __init__(self, *, seq_len, patch_size, num_classes, dim, depth, heads, mlp_dim, channels = 3, dim_head = 64):
@@ -93,10 +94,7 @@ class SimpleViT(nn.Module):
         self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim)
 
         self.to_latent = nn.Identity()
-        self.linear_head = nn.Sequential(
-            nn.LayerNorm(dim),
-            nn.Linear(dim, num_classes)
-        )
+        self.linear_head = nn.Linear(dim, num_classes)
 
     def forward(self, series):
         *_, n, dtype = *series.shape, series.dtype

vit_pytorch/simple_vit_3d.py

@@ -26,7 +26,7 @@ def posemb_sincos_3d(patches, temperature = 10000, dtype = torch.float32):
 
     z = z.flatten()[:, None] * omega[None, :]
     y = y.flatten()[:, None] * omega[None, :]
-    x = x.flatten()[:, None] * omega[None, :] 
+    x = x.flatten()[:, None] * omega[None, :]
 
     pe = torch.cat((x.sin(), x.cos(), y.sin(), y.cos(), z.sin(), z.cos()), dim = 1)
 
@@ -77,6 +77,7 @@ class Attention(nn.Module):
 class Transformer(nn.Module):
     def __init__(self, dim, depth, heads, dim_head, mlp_dim):
         super().__init__()
+        self.norm = nn.LayerNorm(dim)
         self.layers = nn.ModuleList([])
         for _ in range(depth):
             self.layers.append(nn.ModuleList([
@@ -87,13 +88,13 @@ class Transformer(nn.Module):
         for attn, ff in self.layers:
             x = attn(x) + x
             x = ff(x) + x
-        return x
+        return self.norm(x)
 
 class SimpleViT(nn.Module):
     def __init__(self, *, image_size, image_patch_size, frames, frame_patch_size, num_classes, dim, depth, heads, mlp_dim, channels = 3, dim_head = 64):
         super().__init__()
         image_height, image_width = pair(image_size)
-        patch_height, patch_width = pair(image_patch_size)
+        self.patch_size = patch_height, patch_width = pair(image_patch_size)
 
         assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
         assert frames % frame_patch_size == 0, 'Frames must be divisible by the frame patch size'
@@ -102,7 +103,7 @@ class SimpleViT(nn.Module):
         patch_dim = channels * patch_height * patch_width * frame_patch_size
 
         self.to_patch_embedding = nn.Sequential(
-            Rearrange('b c (f pf) (h p1) (w p2) -> b f h w (p1 p2 pf c)', p1 = patch_height, p2 = patch_width, pf = frame_patch_size),
+            Rearrange('b c (f pf) (h p1) (w p2) -> b f h w (pf p1 p2 c)', p1 = patch_height, p2 = patch_width, pf = frame_patch_size),
             nn.LayerNorm(patch_dim),
             nn.Linear(patch_dim, dim),
             nn.LayerNorm(dim),
@@ -111,10 +112,7 @@ class SimpleViT(nn.Module):
         self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim)
 
         self.to_latent = nn.Identity()
-        self.linear_head = nn.Sequential(
-            nn.LayerNorm(dim),
-            nn.Linear(dim, num_classes)
-        )
+        self.linear_head = nn.Linear(dim, num_classes)
 
     def forward(self, video):
         *_, h, w, dtype = *video.shape, video.dtype

vit_pytorch/simple_vit_attn_residual.py

@@ -0,0 +1,243 @@
+# @bycloud explanation https://www.youtube.com/watch?v=iw1VF8HOCrk
+
+from __future__ import annotations
+
+import torch
+from torch import nn, Tensor
+from torch.nn import Module, ModuleList
+
+from einops import rearrange, repeat
+from einops.layers.torch import Rearrange
+
+# helpers
+
+def exists(v):
+    return v is not None
+
+def default(v, d):
+    return v if exists(v) else d
+
+def last(arr):
+    return arr[-1]
+
+def pair(t):
+    return t if isinstance(t, tuple) else (t, t)
+
+def divisible_by(num, den):
+    return (num % den) == 0
+
+def posemb_sincos_2d(h, w, dim, temperature = 10000, dtype = torch.float32):
+    y, x = torch.meshgrid(torch.arange(h), torch.arange(w), indexing = 'ij')
+    assert divisible_by(dim, 4), 'feature dimension must be multiple of 4 for sincos emb'
+    omega = torch.arange(dim // 4) / (dim // 4 - 1)
+    omega = 1.0 / (temperature ** omega)
+
+    y = y.flatten()[:, None] * omega[None, :]
+    x = x.flatten()[:, None] * omega[None, :]
+    pe = torch.cat((x.sin(), x.cos(), y.sin(), y.cos()), dim = 1)
+    return pe.type(dtype)
+
+# classes
+
+class FeedForward(Module):
+    def __init__(self, dim, hidden_dim):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Linear(dim, hidden_dim),
+            nn.GELU(),
+            nn.Linear(hidden_dim, dim),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+class Attention(Module):
+    def __init__(self, dim, heads = 8, dim_head = 64, cross_attend = False):
+        super().__init__()
+        inner_dim = dim_head * heads
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+
+        self.norm = nn.LayerNorm(dim)
+        self.norm_context = nn.LayerNorm(dim) if cross_attend else None
+
+        self.to_q = nn.Linear(dim, inner_dim, bias = False)
+        self.to_kv = nn.Linear(dim, inner_dim * 2, bias = False)
+        self.to_out = nn.Linear(inner_dim, dim, bias = False)
+
+    def forward(self, x, context = None):
+        x = self.norm(x)
+
+        if exists(context):
+            context = self.norm_context(context)
+        else:
+            context = x
+
+        q = self.to_q(x)
+        k, v = self.to_kv(context).chunk(2, dim = -1)
+
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), (q, k, v))
+
+        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+        attn = dots.softmax(dim = -1)
+
+        out = torch.matmul(attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)
+
+class AttentionResidual(Module):
+    def __init__(self, fn, dim, heads = 8, dim_head = 64, learned_query = True, disable = False):
+        super().__init__()
+        self.fn = fn
+        self.disable = disable
+
+        if disable:
+            return
+
+        self.attn = Attention(dim, heads = heads, dim_head = dim_head, cross_attend = True)
+        self.learned_query = nn.Parameter(torch.randn(dim)) if learned_query else None
+
+    def forward(self, history: list[Tensor]) -> Tensor:
+        if self.disable:
+            return self.fn(last(history))
+
+        batch, seq_len = history[0].shape[:2]
+
+        context = torch.stack(history, dim = 2)
+        context = rearrange(context, 'b n l d -> (b n) l d')
+
+        if exists(self.learned_query):
+            q = repeat(self.learned_query, 'd -> (b n) 1 d', b = batch, n = seq_len)
+        else:
+            q = rearrange(last(history), 'b n d -> (b n) 1 d')
+
+        pooled = self.attn(q, context = context)
+        pooled = rearrange(pooled, '(b n) 1 d -> b n d', b = batch, n = seq_len)
+
+        return self.fn(pooled)
+
+class Transformer(Module):
+    def __init__(self, dim, depth, heads, dim_head, mlp_dim, learned_query = True):
+        super().__init__()
+
+        self.layers = ModuleList([])
+        for ind in range(depth):
+            is_first = ind == 0
+
+            self.layers.append(ModuleList([
+                AttentionResidual(Attention(dim, heads = heads, dim_head = dim_head), dim, heads = heads, dim_head = dim_head, learned_query = learned_query, disable = is_first),
+                AttentionResidual(FeedForward(dim, mlp_dim), dim, heads = heads, dim_head = dim_head, learned_query = learned_query),
+            ]))
+
+        self.final_pool = AttentionResidual(nn.LayerNorm(dim), dim, heads = heads, dim_head = dim_head, learned_query = learned_query)
+
+    def forward(
+        self,
+        x,
+        history: list[Tensor] | None = None,
+        return_history = False
+    ):
+        history = [*default(history, [])]
+
+        history.append(x)
+
+        for attn_residual, ff_residual in self.layers:
+            history.append(attn_residual(history))
+            history.append(ff_residual(history))
+
+        out = self.final_pool(history)
+
+        if return_history:
+            return out, history
+
+        return out
+
+class SimpleViTAttnResidual(Module):
+    def __init__(
+        self,
+        *,
+        image_size,
+        patch_size,
+        num_classes,
+        dim,
+        depth,
+        heads,
+        mlp_dim,
+        channels = 3,
+        dim_head = 64,
+        learned_query = True
+    ):
+        super().__init__()
+        image_height, image_width = pair(image_size)
+        self.patch_size = patch_height, patch_width = pair(patch_size)
+
+        assert divisible_by(image_height, patch_height) and divisible_by(image_width, patch_width), 'Image dimensions must be divisible by the patch size.'
+
+        patch_dim = channels * patch_height * patch_width
+
+        self.to_patch_embedding = nn.Sequential(
+            Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_height, p2 = patch_width),
+            nn.LayerNorm(patch_dim),
+            nn.Linear(patch_dim, dim),
+            nn.LayerNorm(dim),
+        )
+
+        self.pos_embedding = posemb_sincos_2d(
+            h = image_height // patch_height,
+            w = image_width // patch_width,
+            dim = dim,
+        )
+
+        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, learned_query = learned_query)
+
+        self.to_latent = nn.Identity()
+        self.linear_head = nn.Linear(dim, num_classes)
+
+    def forward(
+        self,
+        img,
+        history: list[Tensor] | None = None,
+        return_history = False
+    ):
+        device, dtype = img.device, img.dtype
+
+        x = self.to_patch_embedding(img)
+        x += self.pos_embedding.to(device, dtype = dtype)
+
+        x = self.transformer(x, history = history, return_history = return_history)
+
+        if return_history:
+            x, history = x
+
+        x = x.mean(dim = 1)
+
+        x = self.to_latent(x)
+        out = self.linear_head(x)
+
+        if return_history:
+            return out, history
+
+        return out
+
+if __name__ == '__main__':
+    for learned_query in (True, False):
+        v = SimpleViTAttnResidual(
+            image_size = 256,
+            patch_size = 32,
+            num_classes = 1000,
+            dim = 1024,
+            depth = 6,
+            heads = 16,
+            mlp_dim = 2048,
+            learned_query = learned_query
+        )
+
+        img = torch.randn(2, 3, 256, 256)
+        preds, history = v(img, return_history = True)
+
+        assert preds.shape == (2, 1000)
+
+        preds, _ = v(img, history = history, return_history = True)
+
+        assert preds.shape == (2, 1000)

vit_pytorch/simple_vit_orthog_residual_update.py

@@ -0,0 +1,206 @@
+# Revisiting Residual Connections: Orthogonal Updates for Stable and Efficient Deep Networks
+# Giyeong Oh et al. https://arxiv.org/abs/2505.11881
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.nn import Module, ModuleList
+
+from einops import rearrange
+from einops.layers.torch import Rearrange
+
+# helpers
+
+def exists(v):
+    return v is not None
+
+def pair(t):
+    return t if isinstance(t, tuple) else (t, t)
+
+def posemb_sincos_2d(h, w, dim, temperature: int = 10000, dtype = torch.float32):
+    y, x = torch.meshgrid(torch.arange(h), torch.arange(w), indexing="ij")
+    assert (dim % 4) == 0, "feature dimension must be multiple of 4 for sincos emb"
+    omega = torch.arange(dim // 4) / (dim // 4 - 1)
+    omega = 1.0 / (temperature ** omega)
+
+    y = y.flatten()[:, None] * omega[None, :]
+    x = x.flatten()[:, None] * omega[None, :]
+    pe = torch.cat((x.sin(), x.cos(), y.sin(), y.cos()), dim=1)
+    return pe.type(dtype)
+
+# classes
+
+class FeedForward(Module):
+    def __init__(self, dim, hidden_dim):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Linear(dim, hidden_dim),
+            nn.GELU(),
+            nn.Linear(hidden_dim, dim),
+        )
+    def forward(self, x):
+        return self.net(x)
+
+class Attention(Module):
+    def __init__(self, dim, heads = 8, dim_head = 64):
+        super().__init__()
+        inner_dim = dim_head *  heads
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+        self.norm = nn.LayerNorm(dim)
+
+        self.attend = nn.Softmax(dim = -1)
+
+        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
+        self.to_out = nn.Linear(inner_dim, dim, bias = False)
+
+    def forward(self, x):
+        x = self.norm(x)
+
+        qkv = self.to_qkv(x).chunk(3, dim = -1)
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
+
+        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+
+        attn = self.attend(dots)
+
+        out = torch.matmul(attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)
+
+class OrthogonalResidualUpdate(Module):
+    def __init__(
+        self,
+        block: Module,
+        dim = None,
+        double_precision = True,
+        learned = False
+    ):
+        super().__init__()
+        self.block = block
+        self.double_precision = double_precision
+
+        self.learned = learned
+
+        if learned:
+            assert exists(dim)
+            self.to_modulation = nn.Linear(dim, 2)
+
+    def orthog_proj(self, block_out, residual):
+        use_double, dtype = self.double_precision, residual.dtype
+
+        if use_double:
+            residual, block_out = residual.double(), block_out.double()
+
+        # get orthogonal projection of the attention or feedforward output respect to residual
+
+        unit = F.normalize(residual, dim = -1)
+        parallel = (block_out * unit).sum(dim = -1, keepdim = True) * unit
+        orthogonal = block_out - parallel
+
+        # back to original dtype if double precision
+
+        if use_double:
+            parallel, orthogonal = parallel.to(dtype), orthogonal.to(dtype)
+
+        return parallel, orthogonal
+
+    def forward(self, residual):
+        block_out = self.block(residual)
+
+        parallel_update, orthog_update = self.orthog_proj(block_out, residual)
+
+        if self.learned:
+            parallel_mod, orthog_mod = self.to_modulation(block_out).sigmoid().split(1, dim = -1)
+            parallel_update = parallel_update * parallel_mod
+            orthog_update = orthog_update * orthog_mod
+        else:
+            parallel_update = 0
+
+        return residual + parallel_update + orthog_update
+
+class Transformer(Module):
+    def __init__(self, dim, depth, heads, dim_head, mlp_dim, orthog_residual_update_kwargs: dict = dict()):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.layers = ModuleList([])
+
+        for _ in range(depth):
+            attn = Attention(dim, heads = heads, dim_head = dim_head)
+            ff = FeedForward(dim, mlp_dim)
+
+            self.layers.append(ModuleList([
+                OrthogonalResidualUpdate(attn, dim = dim, **orthog_residual_update_kwargs),
+                OrthogonalResidualUpdate(ff, dim = dim, **orthog_residual_update_kwargs)
+            ]))
+
+    def forward(self, x):
+
+        for attn, ff in self.layers:
+            x = attn(x)
+            x = ff(x)
+
+        return self.norm(x)
+
+class SimpleViT(Module):
+    def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, channels = 3, dim_head = 64, orthog_residual_update_kwargs: dict = dict()):
+        super().__init__()
+        image_height, image_width = pair(image_size)
+        self.patch_size = patch_height, patch_width = pair(patch_size)
+
+        assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
+
+        patch_dim = channels * patch_height * patch_width
+
+        self.to_patch_embedding = nn.Sequential(
+            Rearrange("b c (h p1) (w p2) -> b (h w) (p1 p2 c)", p1 = patch_height, p2 = patch_width),
+            nn.LayerNorm(patch_dim),
+            nn.Linear(patch_dim, dim),
+            nn.LayerNorm(dim),
+        )
+
+        self.pos_embedding = posemb_sincos_2d(
+            h = image_height // patch_height,
+            w = image_width // patch_width,
+            dim = dim,
+        )
+
+        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, orthog_residual_update_kwargs)
+
+        self.pool = "mean"
+        self.to_latent = nn.Identity()
+
+        self.linear_head = nn.Linear(dim, num_classes)
+
+    def forward(self, img):
+        device = img.device
+
+        x = self.to_patch_embedding(img)
+        x += self.pos_embedding.to(device, dtype=x.dtype)
+
+        x = self.transformer(x)
+        x = x.mean(dim = 1)
+
+        x = self.to_latent(x)
+        return self.linear_head(x)
+
+# quick test
+
+if __name__ == '__main__':
+    vit = SimpleViT(
+        image_size = 256,
+        patch_size = 16,
+        num_classes = 10,
+        dim = 512,
+        depth = 2,
+        heads = 4,
+        mlp_dim = 2048,
+        orthog_residual_update_kwargs = dict(
+            learned = True
+        )
+    )
+
+    images = torch.randn(2, 3, 256, 256)
+
+    assert vit(images).shape == (2, 10)

vit_pytorch/simple_vit_with_fft.py

@@ -0,0 +1,162 @@
+import torch
+from torch.fft import fft2
+from torch import nn
+
+from einops import rearrange, reduce, pack, unpack
+from einops.layers.torch import Rearrange
+
+# helpers
+
+def pair(t):
+    return t if isinstance(t, tuple) else (t, t)
+
+def posemb_sincos_2d(h, w, dim, temperature: int = 10000, dtype = torch.float32):
+    y, x = torch.meshgrid(torch.arange(h), torch.arange(w), indexing="ij")
+    assert (dim % 4) == 0, "feature dimension must be multiple of 4 for sincos emb"
+    omega = torch.arange(dim // 4) / (dim // 4 - 1)
+    omega = 1.0 / (temperature ** omega)
+
+    y = y.flatten()[:, None] * omega[None, :]
+    x = x.flatten()[:, None] * omega[None, :]
+    pe = torch.cat((x.sin(), x.cos(), y.sin(), y.cos()), dim=1)
+    return pe.type(dtype)
+
+# classes
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, hidden_dim):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Linear(dim, hidden_dim),
+            nn.GELU(),
+            nn.Linear(hidden_dim, dim),
+        )
+    def forward(self, x):
+        return self.net(x)
+
+class Attention(nn.Module):
+    def __init__(self, dim, heads = 8, dim_head = 64):
+        super().__init__()
+        inner_dim = dim_head *  heads
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+        self.norm = nn.LayerNorm(dim)
+
+        self.attend = nn.Softmax(dim = -1)
+
+        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
+        self.to_out = nn.Linear(inner_dim, dim, bias = False)
+
+    def forward(self, x):
+        x = self.norm(x)
+
+        qkv = self.to_qkv(x).chunk(3, dim = -1)
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
+
+        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+
+        attn = self.attend(dots)
+
+        out = torch.matmul(attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)
+
+class Transformer(nn.Module):
+    def __init__(self, dim, depth, heads, dim_head, mlp_dim):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.layers = nn.ModuleList([])
+        for _ in range(depth):
+            self.layers.append(nn.ModuleList([
+                Attention(dim, heads = heads, dim_head = dim_head),
+                FeedForward(dim, mlp_dim)
+            ]))
+    def forward(self, x):
+        for attn, ff in self.layers:
+            x = attn(x) + x
+            x = ff(x) + x
+        return self.norm(x)
+
+class SimpleViT(nn.Module):
+    def __init__(self, *, image_size, patch_size, freq_patch_size, num_classes, dim, depth, heads, mlp_dim, channels = 3, dim_head = 64):
+        super().__init__()
+        image_height, image_width = pair(image_size)
+        self.patch_size = patch_height, patch_width = pair(patch_size)
+        freq_patch_height, freq_patch_width = pair(freq_patch_size)
+
+        assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
+        assert image_height % freq_patch_height == 0 and image_width % freq_patch_width == 0, 'Image dimensions must be divisible by the freq patch size.'
+
+        patch_dim = channels * patch_height * patch_width
+        freq_patch_dim = channels * 2 * freq_patch_height * freq_patch_width
+
+        self.to_patch_embedding = nn.Sequential(
+            Rearrange("b c (h p1) (w p2) -> b (h w) (p1 p2 c)", p1 = patch_height, p2 = patch_width),
+            nn.LayerNorm(patch_dim),
+            nn.Linear(patch_dim, dim),
+            nn.LayerNorm(dim),
+        )
+
+        self.to_freq_embedding = nn.Sequential(
+            Rearrange("b c (h p1) (w p2) ri -> b (h w) (p1 p2 ri c)", p1 = freq_patch_height, p2 = freq_patch_width),
+            nn.LayerNorm(freq_patch_dim),
+            nn.Linear(freq_patch_dim, dim),
+            nn.LayerNorm(dim)
+        )
+
+        self.pos_embedding = posemb_sincos_2d(
+            h = image_height // patch_height,
+            w = image_width // patch_width,
+            dim = dim,
+        )
+
+        self.freq_pos_embedding = posemb_sincos_2d(
+            h = image_height // freq_patch_height,
+            w = image_width // freq_patch_width,
+            dim = dim
+        )
+
+        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim)
+
+        self.pool = "mean"
+        self.to_latent = nn.Identity()
+
+        self.linear_head = nn.Linear(dim, num_classes)
+
+    def forward(self, img):
+        device, dtype = img.device, img.dtype
+
+        x = self.to_patch_embedding(img)
+        freqs = torch.view_as_real(fft2(img))
+
+        f = self.to_freq_embedding(freqs)
+
+        x += self.pos_embedding.to(device, dtype = dtype)
+        f += self.freq_pos_embedding.to(device, dtype = dtype)
+
+        x, ps = pack((f, x), 'b * d')
+
+        x = self.transformer(x)
+
+        _, x = unpack(x, ps, 'b * d')
+        x = reduce(x, 'b n d -> b d', 'mean')
+
+        x = self.to_latent(x)
+        return self.linear_head(x)
+
+if __name__ == '__main__':
+    vit = SimpleViT(
+        num_classes = 1000,
+        image_size = 256,
+        patch_size = 8,
+        freq_patch_size = 8,
+        dim = 1024,
+        depth = 1,
+        heads = 8,
+        mlp_dim = 2048,
+    )
+
+    images = torch.randn(8, 3, 256, 256)
+
+    logits = vit(images)

vit_pytorch/simple_vit_with_hyper_connections.py

@@ -0,0 +1,233 @@
+"""
+ViT + Hyper-Connections + Register Tokens
+https://arxiv.org/abs/2409.19606
+"""
+
+import torch
+from torch import nn, tensor
+from torch.nn import Module, ModuleList
+
+from einops import rearrange, repeat, reduce, einsum, pack, unpack
+from einops.layers.torch import Rearrange
+
+# b - batch, h - heads, n - sequence, e - expansion rate / residual streams, d - feature dimension
+
+# helpers
+
+def pair(t):
+    return t if isinstance(t, tuple) else (t, t)
+
+def posemb_sincos_2d(h, w, dim, temperature: int = 10000, dtype = torch.float32):
+    y, x = torch.meshgrid(torch.arange(h), torch.arange(w), indexing="ij")
+    assert (dim % 4) == 0, "feature dimension must be multiple of 4 for sincos emb"
+    omega = torch.arange(dim // 4) / (dim // 4 - 1)
+    omega = 1.0 / (temperature ** omega)
+
+    y = y.flatten()[:, None] * omega[None, :]
+    x = x.flatten()[:, None] * omega[None, :]
+    pe = torch.cat((x.sin(), x.cos(), y.sin(), y.cos()), dim=1)
+    return pe.type(dtype)
+
+# hyper connections
+
+class HyperConnection(Module):
+    def __init__(
+        self,
+        dim,
+        num_residual_streams,
+        layer_index
+    ):
+        """ Appendix J - Algorithm 2, Dynamic only """
+        super().__init__()
+
+        self.norm = nn.LayerNorm(dim, bias = False)
+
+        self.num_residual_streams = num_residual_streams
+        self.layer_index = layer_index
+
+        self.static_beta = nn.Parameter(torch.ones(num_residual_streams))
+
+        init_alpha0 = torch.zeros((num_residual_streams, 1))
+        init_alpha0[layer_index % num_residual_streams, 0] = 1.
+
+        self.static_alpha = nn.Parameter(torch.cat([init_alpha0, torch.eye(num_residual_streams)], dim = 1))
+
+        self.dynamic_alpha_fn = nn.Parameter(torch.zeros(dim, num_residual_streams + 1))
+        self.dynamic_alpha_scale = nn.Parameter(tensor(1e-2))
+        self.dynamic_beta_fn = nn.Parameter(torch.zeros(dim))
+        self.dynamic_beta_scale = nn.Parameter(tensor(1e-2))
+
+    def width_connection(self, residuals):
+        normed = self.norm(residuals)
+
+        wc_weight = (normed @ self.dynamic_alpha_fn).tanh()
+        dynamic_alpha = wc_weight * self.dynamic_alpha_scale
+        alpha = dynamic_alpha + self.static_alpha
+
+        dc_weight = (normed @ self.dynamic_beta_fn).tanh()
+        dynamic_beta = dc_weight * self.dynamic_beta_scale
+        beta = dynamic_beta + self.static_beta
+
+        # width connection
+        mix_h = einsum(alpha, residuals, '... e1 e2, ... e1 d -> ... e2 d')
+
+        branch_input, residuals = mix_h[..., 0, :], mix_h[..., 1:, :]
+
+        return branch_input, residuals, beta
+
+    def depth_connection(
+        self,
+        branch_output,
+        residuals,
+        beta
+    ):
+        return einsum(branch_output, beta, "b n d, b n e -> b n e d") + residuals
+
+# classes
+
+class FeedForward(Module):
+    def __init__(self, dim, hidden_dim):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Linear(dim, hidden_dim),
+            nn.GELU(),
+            nn.Linear(hidden_dim, dim),
+        )
+    def forward(self, x):
+        return self.net(x)
+
+class Attention(Module):
+    def __init__(self, dim, heads = 8, dim_head = 64):
+        super().__init__()
+        inner_dim = dim_head *  heads
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+        self.norm = nn.LayerNorm(dim)
+
+        self.attend = nn.Softmax(dim = -1)
+
+        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
+        self.to_out = nn.Linear(inner_dim, dim, bias = False)
+
+    def forward(self, x):
+        x = self.norm(x)
+
+        qkv = self.to_qkv(x).chunk(3, dim = -1)
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
+
+        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+
+        attn = self.attend(dots)
+
+        out = torch.matmul(attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)
+
+class Transformer(Module):
+    def __init__(self, dim, depth, heads, dim_head, mlp_dim, num_residual_streams):
+        super().__init__()
+
+        self.num_residual_streams = num_residual_streams
+
+        self.norm = nn.LayerNorm(dim)
+        self.layers = ModuleList([])
+
+        for layer_index in range(depth):
+            self.layers.append(nn.ModuleList([
+                HyperConnection(dim, num_residual_streams, layer_index),
+                Attention(dim, heads = heads, dim_head = dim_head),
+                HyperConnection(dim, num_residual_streams, layer_index),
+                FeedForward(dim, mlp_dim)
+            ]))
+
+    def forward(self, x):
+
+        x = repeat(x, 'b n d -> b n e d', e = self.num_residual_streams)
+
+        for attn_hyper_conn, attn, ff_hyper_conn, ff in self.layers:
+
+            x, attn_res, beta = attn_hyper_conn.width_connection(x)
+
+            x = attn(x)
+
+            x = attn_hyper_conn.depth_connection(x, attn_res, beta)
+
+            x, ff_res, beta = ff_hyper_conn.width_connection(x)
+
+            x = ff(x)
+
+            x = ff_hyper_conn.depth_connection(x, ff_res, beta)
+
+        x = reduce(x, 'b n e d -> b n d', 'sum')
+
+        return self.norm(x)
+
+class SimpleViT(nn.Module):
+    def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, num_residual_streams, num_register_tokens = 4, channels = 3, dim_head = 64):
+        super().__init__()
+        image_height, image_width = pair(image_size)
+        self.patch_size = patch_height, patch_width = pair(patch_size)
+
+        assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
+
+        patch_dim = channels * patch_height * patch_width
+
+        self.to_patch_embedding = nn.Sequential(
+            Rearrange("b c (h p1) (w p2) -> b (h w) (p1 p2 c)", p1 = patch_height, p2 = patch_width),
+            nn.LayerNorm(patch_dim),
+            nn.Linear(patch_dim, dim),
+            nn.LayerNorm(dim),
+        )
+
+        self.register_tokens = nn.Parameter(torch.randn(num_register_tokens, dim))
+
+        self.pos_embedding = posemb_sincos_2d(
+            h = image_height // patch_height,
+            w = image_width // patch_width,
+            dim = dim,
+        )
+
+        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, num_residual_streams)
+
+        self.pool = "mean"
+        self.to_latent = nn.Identity()
+
+        self.linear_head = nn.Linear(dim, num_classes)
+
+    def forward(self, img):
+        batch, device = img.shape[0], img.device
+
+        x = self.to_patch_embedding(img)
+        x += self.pos_embedding.to(x)
+
+        r = repeat(self.register_tokens, 'n d -> b n d', b = batch)
+
+        x, ps = pack([x, r], 'b * d')
+
+        x = self.transformer(x)
+
+        x, _ = unpack(x, ps, 'b * d')
+
+        x = x.mean(dim = 1)
+
+        x = self.to_latent(x)
+        return self.linear_head(x)
+
+# main
+
+if __name__ == '__main__':
+    vit = SimpleViT(
+        num_classes = 1000,
+        image_size = 256,
+        patch_size = 8,
+        dim = 1024,
+        depth = 12,
+        heads = 8,
+        mlp_dim = 2048,
+        num_residual_streams = 8
+    )
+
+    images = torch.randn(3, 3, 256, 256)
+
+    logits = vit(images)

vit_pytorch/simple_vit_with_patch_dropout.py

@@ -18,7 +18,7 @@ def posemb_sincos_2d(patches, temperature = 10000, dtype = torch.float32):
     omega = 1. / (temperature ** omega)
 
     y = y.flatten()[:, None] * omega[None, :]
-    x = x.flatten()[:, None] * omega[None, :] 
+    x = x.flatten()[:, None] * omega[None, :]
     pe = torch.cat((x.sin(), x.cos(), y.sin(), y.cos()), dim = 1)
     return pe.type(dtype)
 
@@ -87,6 +87,7 @@ class Attention(nn.Module):
 class Transformer(nn.Module):
     def __init__(self, dim, depth, heads, dim_head, mlp_dim):
         super().__init__()
+        self.norm = nn.LayerNorm(dim)
         self.layers = nn.ModuleList([])
         for _ in range(depth):
             self.layers.append(nn.ModuleList([
@@ -97,13 +98,13 @@ class Transformer(nn.Module):
         for attn, ff in self.layers:
             x = attn(x) + x
             x = ff(x) + x
-        return x
+        return self.norm(x)
 
 class SimpleViT(nn.Module):
     def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, channels = 3, dim_head = 64, patch_dropout = 0.5):
         super().__init__()
         image_height, image_width = pair(image_size)
-        patch_height, patch_width = pair(patch_size)
+        self.patch_size = patch_height, patch_width = pair(patch_size)
 
         assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
 
@@ -122,10 +123,7 @@ class SimpleViT(nn.Module):
         self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim)
 
         self.to_latent = nn.Identity()
-        self.linear_head = nn.Sequential(
-            nn.LayerNorm(dim),
-            nn.Linear(dim, num_classes)
-        )
+        self.linear_head = nn.Linear(dim, num_classes)
 
     def forward(self, img):
         *_, h, w, dtype = *img.shape, img.dtype

vit_pytorch/simple_vit_with_qk_norm.py

@@ -0,0 +1,141 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from einops import rearrange
+from einops.layers.torch import Rearrange
+
+# helpers
+
+def pair(t):
+    return t if isinstance(t, tuple) else (t, t)
+
+def posemb_sincos_2d(h, w, dim, temperature: int = 10000, dtype = torch.float32):
+    y, x = torch.meshgrid(torch.arange(h), torch.arange(w), indexing="ij")
+    assert (dim % 4) == 0, "feature dimension must be multiple of 4 for sincos emb"
+    omega = torch.arange(dim // 4) / (dim // 4 - 1)
+    omega = 1.0 / (temperature ** omega)
+
+    y = y.flatten()[:, None] * omega[None, :]
+    x = x.flatten()[:, None] * omega[None, :]
+    pe = torch.cat((x.sin(), x.cos(), y.sin(), y.cos()), dim=1)
+    return pe.type(dtype)
+
+# they use a query-key normalization that is equivalent to rms norm (no mean-centering, learned gamma), from vit 22B paper
+
+# in latest tweet, seem to claim more stable training at higher learning rates
+# unsure if this has taken off within Brain, or it has some hidden drawback
+
+class RMSNorm(nn.Module):
+    def __init__(self, heads, dim):
+        super().__init__()
+        self.scale = dim ** 0.5
+        self.gamma = nn.Parameter(torch.ones(heads, 1, dim) / self.scale)
+
+    def forward(self, x):
+        normed = F.normalize(x, dim = -1)
+        return normed * self.scale * self.gamma
+
+# classes
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, hidden_dim):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Linear(dim, hidden_dim),
+            nn.GELU(),
+            nn.Linear(hidden_dim, dim),
+        )
+    def forward(self, x):
+        return self.net(x)
+
+class Attention(nn.Module):
+    def __init__(self, dim, heads = 8, dim_head = 64):
+        super().__init__()
+        inner_dim = dim_head *  heads
+        self.heads = heads
+        self.norm = nn.LayerNorm(dim)
+
+        self.attend = nn.Softmax(dim = -1)
+
+        self.q_norm = RMSNorm(heads, dim_head)
+        self.k_norm = RMSNorm(heads, dim_head)
+
+        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
+        self.to_out = nn.Linear(inner_dim, dim, bias = False)
+
+    def forward(self, x):
+        x = self.norm(x)
+
+        qkv = self.to_qkv(x).chunk(3, dim = -1)
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
+
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+
+        dots = torch.matmul(q, k.transpose(-1, -2))
+
+        attn = self.attend(dots)
+
+        out = torch.matmul(attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)
+
+class Transformer(nn.Module):
+    def __init__(self, dim, depth, heads, dim_head, mlp_dim):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.layers = nn.ModuleList([])
+        for _ in range(depth):
+            self.layers.append(nn.ModuleList([
+                Attention(dim, heads = heads, dim_head = dim_head),
+                FeedForward(dim, mlp_dim)
+            ]))
+    def forward(self, x):
+        for attn, ff in self.layers:
+            x = attn(x) + x
+            x = ff(x) + x
+        return self.norm(x)
+
+class SimpleViT(nn.Module):
+    def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, channels = 3, dim_head = 64):
+        super().__init__()
+        image_height, image_width = pair(image_size)
+        self.patch_size = patch_height, patch_width = pair(patch_size)
+
+        assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
+
+        patch_dim = channels * patch_height * patch_width
+
+        self.to_patch_embedding = nn.Sequential(
+            Rearrange("b c (h p1) (w p2) -> b (h w) (p1 p2 c)", p1 = patch_height, p2 = patch_width),
+            nn.LayerNorm(patch_dim),
+            nn.Linear(patch_dim, dim),
+            nn.LayerNorm(dim),
+        )
+
+        self.pos_embedding = posemb_sincos_2d(
+            h = image_height // patch_height,
+            w = image_width // patch_width,
+            dim = dim,
+        )
+
+        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim)
+
+        self.pool = "mean"
+        self.to_latent = nn.Identity()
+
+        self.linear_head = nn.LayerNorm(dim)
+
+    def forward(self, img):
+        device = img.device
+
+        x = self.to_patch_embedding(img)
+        x += self.pos_embedding.to(device, dtype=x.dtype)
+
+        x = self.transformer(x)
+        x = x.mean(dim = 1)
+
+        x = self.to_latent(x)
+        return self.linear_head(x)

vit_pytorch/simple_vit_with_register_tokens.py

@@ -0,0 +1,134 @@
+"""
+    Vision Transformers Need Registers
+    https://arxiv.org/abs/2309.16588
+"""
+
+import torch
+from torch import nn
+
+from einops import rearrange, repeat, pack, unpack
+from einops.layers.torch import Rearrange
+
+# helpers
+
+def pair(t):
+    return t if isinstance(t, tuple) else (t, t)
+
+def posemb_sincos_2d(h, w, dim, temperature: int = 10000, dtype = torch.float32):
+    y, x = torch.meshgrid(torch.arange(h), torch.arange(w), indexing="ij")
+    assert (dim % 4) == 0, "feature dimension must be multiple of 4 for sincos emb"
+    omega = torch.arange(dim // 4) / (dim // 4 - 1)
+    omega = 1.0 / (temperature ** omega)
+
+    y = y.flatten()[:, None] * omega[None, :]
+    x = x.flatten()[:, None] * omega[None, :]
+    pe = torch.cat((x.sin(), x.cos(), y.sin(), y.cos()), dim=1)
+    return pe.type(dtype)
+
+# classes
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, hidden_dim):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Linear(dim, hidden_dim),
+            nn.GELU(),
+            nn.Linear(hidden_dim, dim),
+        )
+    def forward(self, x):
+        return self.net(x)
+
+class Attention(nn.Module):
+    def __init__(self, dim, heads = 8, dim_head = 64):
+        super().__init__()
+        inner_dim = dim_head *  heads
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+        self.norm = nn.LayerNorm(dim)
+
+        self.attend = nn.Softmax(dim = -1)
+
+        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
+        self.to_out = nn.Linear(inner_dim, dim, bias = False)
+
+    def forward(self, x):
+        x = self.norm(x)
+
+        qkv = self.to_qkv(x).chunk(3, dim = -1)
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
+
+        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+
+        attn = self.attend(dots)
+
+        out = torch.matmul(attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)
+
+class Transformer(nn.Module):
+    def __init__(self, dim, depth, heads, dim_head, mlp_dim):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.layers = nn.ModuleList([])
+        for _ in range(depth):
+            self.layers.append(nn.ModuleList([
+                Attention(dim, heads = heads, dim_head = dim_head),
+                FeedForward(dim, mlp_dim)
+            ]))
+    def forward(self, x):
+        for attn, ff in self.layers:
+            x = attn(x) + x
+            x = ff(x) + x
+        return self.norm(x)
+
+class SimpleViT(nn.Module):
+    def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, num_register_tokens = 4, channels = 3, dim_head = 64):
+        super().__init__()
+        image_height, image_width = pair(image_size)
+        self.patch_size = patch_height, patch_width = pair(patch_size)
+
+        assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
+
+        patch_dim = channels * patch_height * patch_width
+
+        self.to_patch_embedding = nn.Sequential(
+            Rearrange("b c (h p1) (w p2) -> b (h w) (p1 p2 c)", p1 = patch_height, p2 = patch_width),
+            nn.LayerNorm(patch_dim),
+            nn.Linear(patch_dim, dim),
+            nn.LayerNorm(dim),
+        )
+
+        self.register_tokens = nn.Parameter(torch.randn(num_register_tokens, dim))
+
+        self.pos_embedding = posemb_sincos_2d(
+            h = image_height // patch_height,
+            w = image_width // patch_width,
+            dim = dim,
+        )
+
+        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim)
+
+        self.pool = "mean"
+        self.to_latent = nn.Identity()
+
+        self.linear_head = nn.Linear(dim, num_classes)
+
+    def forward(self, img):
+        batch, device = img.shape[0], img.device
+
+        x = self.to_patch_embedding(img)
+        x += self.pos_embedding.to(device, dtype=x.dtype)
+
+        r = repeat(self.register_tokens, 'n d -> b n d', b = batch)
+
+        x, ps = pack([x, r], 'b * d')
+
+        x = self.transformer(x)
+
+        x, _ = unpack(x, ps, 'b * d')
+
+        x = x.mean(dim = 1)
+
+        x = self.to_latent(x)
+        return self.linear_head(x)

vit_pytorch/simple_vit_with_sparse_gating.py

@@ -0,0 +1,263 @@
+import torch
+from torch import cat
+from torch.nn import Module, ModuleList, Sequential, Linear, LayerNorm, GELU, Softmax, Identity, Parameter
+from torch.nn.functional import pad
+
+from einops import rearrange, repeat, pack, unpack, reduce, einsum
+from einops.layers.torch import Rearrange
+
+# helpers
+
+def pair(t):
+    return t if isinstance(t, tuple) else (t, t)
+
+def log(t, eps = 1e-20):
+    return t.clamp(min = eps).log()
+
+def divisible_by(num, den):
+    return (num % den) == 0
+
+def posemb_sincos_2d(h, w, dim, temperature = 10000, dtype = torch.float32):
+    y, x = torch.meshgrid(torch.arange(h), torch.arange(w), indexing = 'ij')
+    assert divisible_by(dim, 4), 'feature dimension must be multiple of 4 for sincos emb'
+
+    omega = torch.arange(dim // 4) / (dim // 4 - 1)
+    omega = temperature ** -omega
+
+    y, x = y.flatten(), x.flatten()
+
+    y = torch.outer(y, omega)
+    x = torch.outer(x, omega)
+
+    pe = cat((x.sin(), x.cos(), y.sin(), y.cos()), dim = -1)
+    return pe.type(dtype)
+
+# classes
+
+class FeedForward(Module):
+    def __init__(self, dim, hidden_dim):
+        super().__init__()
+        self.net = Sequential(
+            LayerNorm(dim),
+            Linear(dim, hidden_dim),
+            GELU(),
+            Linear(hidden_dim, dim),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+class Attention(Module):
+    def __init__(self, dim, heads = 8, dim_head = 64, dim_mask = 4, num_register_tokens = 4):
+        super().__init__()
+        inner_dim = dim_head *  heads
+        mask_inner_dim = dim_mask * heads
+
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+        self.mask_scale = dim_mask ** -0.5
+        self.num_register_tokens = num_register_tokens
+
+        self.norm = LayerNorm(dim)
+        self.attend = Softmax(dim = -1)
+
+        self.split_dims = (inner_dim, inner_dim, inner_dim, mask_inner_dim, mask_inner_dim)
+
+        self.to_qkv_and_mask = Linear(dim, sum(self.split_dims), bias = False)
+
+        self.to_out = Linear(inner_dim, dim, bias = False)
+
+    def forward(self, x):
+        x = self.norm(x)
+
+        q, k, v, q_mask, k_mask = self.to_qkv_and_mask(x).split(self.split_dims, dim = -1)
+
+        q, k, v, q_mask, k_mask = (rearrange(t, 'b n (h d) -> b h n d', h = self.heads) for t in (q, k, v, q_mask, k_mask))
+
+        # slice out registers
+
+        k_mask_patches = k_mask[:, :, self.num_register_tokens:]
+
+        mask_sim = einsum(q_mask, k_mask_patches, 'b h i d, b h j d -> b h i j') * self.mask_scale
+
+        mask_prob = mask_sim.sigmoid()
+        mask_log = log(mask_prob)
+
+        # pad registers with 0s on left so they act as attention sinks
+
+        mask_log = pad(mask_log, (self.num_register_tokens, 0), value = 0.)
+
+        sim = einsum(q, k, 'b h i d, b h j d -> b h i j') * self.scale
+
+        # add log of mask probabilities to similarities
+
+        sim = sim + mask_log
+
+        attn = self.attend(sim)
+
+        out = einsum(attn, v, 'b h i j, b h j d -> b h i d')
+        out = rearrange(out, 'b h n d -> b n (h d)')
+
+        return self.to_out(out), mask_prob
+
+class Transformer(Module):
+    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dim_mask, num_register_tokens):
+        super().__init__()
+        self.norm = LayerNorm(dim)
+
+        self.layers = ModuleList([ModuleList([
+            Attention(dim, heads = heads, dim_head = dim_head, dim_mask = dim_mask, num_register_tokens = num_register_tokens),
+            FeedForward(dim, mlp_dim)
+        ]) for _ in range(depth)])
+
+    def forward(self, x):
+        mask_probs = []
+
+        for attn, ff in self.layers:
+            attn_out, mask_prob = attn(x)
+
+            x = attn_out + x
+            x = ff(x) + x
+
+            mask_probs.append(mask_prob)
+
+        return self.norm(x), mask_probs
+
+class SimpleViTWithSparseGating(Module):
+    def __init__(
+        self,
+        *,
+        image_size,
+        patch_size,
+        num_classes,
+        dim,
+        depth,
+        heads,
+        mlp_dim,
+        num_register_tokens = 4,
+        channels = 3,
+        dim_head = 64,
+        dim_mask = 4,
+        sparsity_loss_type = 'thresholded_l1',
+        sparsity_threshold = 0.05
+    ):
+        super().__init__()
+        image_height, image_width = pair(image_size)
+        self.patch_size = patch_height, patch_width = pair(patch_size)
+
+        assert divisible_by(image_height, patch_height) and divisible_by(image_width, patch_width), 'Image dimensions must be divisible by the patch size.'
+
+        patch_dim = channels * patch_height * patch_width
+
+        self.to_patch_embedding = Sequential(
+            Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_height, p2 = patch_width),
+            LayerNorm(patch_dim),
+            Linear(patch_dim, dim),
+            LayerNorm(dim),
+        )
+
+        self.num_register_tokens = num_register_tokens
+        self.register_tokens = Parameter(torch.randn(num_register_tokens, dim))
+
+        self.pos_embedding = posemb_sincos_2d(
+            h = image_height // patch_height,
+            w = image_width // patch_width,
+            dim = dim,
+        )
+
+        self.sparsity_loss_type = sparsity_loss_type
+        self.sparsity_threshold = sparsity_threshold
+
+        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dim_mask, num_register_tokens)
+
+        self.pool = 'mean'
+        self.to_latent = Identity()
+
+        self.linear_head = Linear(dim, num_classes)
+
+    def forward(self, img, return_loss = False):
+        batch, device = img.shape[0], img.device
+
+        x = self.to_patch_embedding(img)
+        x += self.pos_embedding.to(device, dtype = x.dtype)
+
+        r = repeat(self.register_tokens, 'n d -> b n d', b = batch)
+
+        # prepend registers
+
+        x, ps = pack([r, x], 'b * d')
+
+        # attend
+
+        x, mask_probs = self.transformer(x)
+
+        # discard registers
+
+        _, x = unpack(x, ps, 'b * d')
+
+        # pool patches
+
+        x = reduce(x, 'b n d -> b d', 'mean')
+
+        # to logits
+
+        x = self.to_latent(x)
+        logits = self.linear_head(x)
+
+        if not return_loss:
+            return logits
+
+        # calculate sparsity loss per image
+
+        p_mean = reduce(mask_probs, 'd b h i j -> b', 'mean')
+
+        if self.sparsity_loss_type == 'l1':
+            loss = p_mean.mean()
+        elif self.sparsity_loss_type == 'thresholded_l1':
+            loss = (p_mean - self.sparsity_threshold).relu().mean()
+        else:
+            raise ValueError(f"Unknown sparsity loss type: {self.sparsity_loss_type}")
+
+        return logits, loss
+
+if __name__ == '__main__':
+    v = SimpleViTWithSparseGating(
+        image_size = 256,
+        patch_size = 32,
+        num_classes = 1000,
+        dim = 1024,
+        depth = 6,
+        heads = 16,
+        mlp_dim = 2048,
+        num_register_tokens = 4,
+        dim_mask = 4,
+        sparsity_loss_type = 'thresholded_l1',
+        sparsity_threshold = 0.05
+    )
+
+    img = torch.randn(2, 3, 256, 256)
+
+    logits = v(img)
+    assert logits.shape == (2, 1000)
+
+    logits, loss = v(img, return_loss = True)
+    assert logits.shape == (2, 1000)
+    assert loss.ndim == 0
+
+    loss.backward()
+
+    v_l1 = SimpleViTWithSparseGating(
+        image_size = 256,
+        patch_size = 32,
+        num_classes = 1000,
+        dim = 1024,
+        depth = 6,
+        heads = 16,
+        mlp_dim = 2048,
+        num_register_tokens = 4,
+        dim_mask = 4,
+        sparsity_loss_type = 'l1'
+    )
+
+    logits, loss_l1 = v_l1(img, return_loss = True)
+    loss_l1.backward()

vit_pytorch/simple_vit_with_specialized_cls.py

@@ -0,0 +1,205 @@
+from __future__ import annotations
+
+# Alexis Marouani et al. https://arxiv.org/abs/2602.08626
+
+import torch
+from torch import nn, cat, Tensor, is_tensor
+from torch.nn import Module, ModuleList
+
+from einops import rearrange, repeat
+from einops.layers.torch import Rearrange
+
+# helpers
+
+def exists(v):
+    return v is not None
+
+def default(v, d):
+    return v if exists(v) else d
+
+def pair(t):
+    return t if isinstance(t, tuple) else (t, t)
+
+def posemb_sincos_2d(h, w, dim, temperature: int = 10000, dtype = torch.float32):
+    y, x = torch.meshgrid(torch.arange(h), torch.arange(w), indexing="ij")
+    assert (dim % 4) == 0, "feature dimension must be multiple of 4 for sincos emb"
+    omega = torch.arange(dim // 4) / (dim // 4 - 1)
+    omega = 1.0 / (temperature ** omega)
+
+    y = y.flatten()[:, None] * omega[None, :]
+    x = x.flatten()[:, None] * omega[None, :]
+    pe = cat((x.sin(), x.cos(), y.sin(), y.cos()), dim=1)
+    return pe.type(dtype)
+
+# classes
+
+class Specialized(Module):
+    def __init__(self, modules: list[Module]):
+        super().__init__()
+        self.fns = ModuleList(modules)
+
+    def forward(
+        self,
+        x: Tensor | list[Tensor],
+        token_lens: tuple[int, ...] = None
+    ):
+        if is_tensor(x):
+            assert exists(token_lens)
+            x = x.split(token_lens, dim = 1)
+
+        assert len(self.fns) == len(x)
+
+        out = tuple(fn(t) for fn, t in zip(self.fns, x))
+
+        if is_tensor:
+            out = cat(out, dim = 1)
+
+        return out
+
+class FeedForward(Module):
+    def __init__(self, dim, hidden_dim):
+        super().__init__()
+        self.norm = Specialized([
+            nn.LayerNorm(dim),
+            nn.LayerNorm(dim)
+        ])
+
+        self.net = nn.Sequential(
+            nn.Linear(dim, hidden_dim),
+            nn.GELU(),
+            nn.Linear(hidden_dim, dim),
+        )
+    def forward(self, x, token_lens = None):
+        x = self.norm(x, token_lens = token_lens)
+        return self.net(x)
+
+class Attention(Module):
+    def __init__(self, dim, heads = 8, dim_head = 64, specialize_qkv = False):
+        super().__init__()
+        inner_dim = dim_head *  heads
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+
+        self.norm = Specialized([
+            nn.LayerNorm(dim),
+            nn.LayerNorm(dim)
+        ])
+
+        self.attend = nn.Softmax(dim = -1)
+        self.specialize_qkv = specialize_qkv
+
+        if specialize_qkv:
+            self.to_qkv = Specialized([
+                nn.Linear(dim, inner_dim * 3, bias = False),
+                nn.Linear(dim, inner_dim * 3, bias = False)
+            ])
+        else:
+            self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
+
+        self.to_out = nn.Linear(inner_dim, dim, bias = False)
+
+    def forward(self, x, token_lens = None):
+        x = self.norm(x, token_lens = token_lens)
+
+        if self.specialize_qkv:
+            qkv = self.to_qkv(x, token_lens = token_lens).chunk(3, dim = -1)
+        else:
+            qkv = self.to_qkv(x).chunk(3, dim = -1)
+
+        q, k, v = (rearrange(t, 'b n (h d) -> b h n d', h = self.heads) for t in qkv)
+
+        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+
+        attn = self.attend(dots)
+
+        out = torch.matmul(attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)
+
+class Transformer(Module):
+    def __init__(self, dim, depth, heads, dim_head, mlp_dim, specialize_qkv_depth):
+        super().__init__()
+        self.norm = Specialized([nn.LayerNorm(dim), nn.LayerNorm(dim)])
+
+        self.layers = ModuleList([])
+
+        for ind in range(depth):
+            specialize_qkv = ind < specialize_qkv_depth
+            self.layers.append(ModuleList([
+                Attention(dim, heads = heads, dim_head = dim_head, specialize_qkv = specialize_qkv),
+                FeedForward(dim, mlp_dim)
+            ]))
+
+    def forward(self, x, token_lens = None):
+        for attn, ff in self.layers:
+            x = attn(x, token_lens = token_lens) + x
+            x = ff(x, token_lens = token_lens) + x
+
+        return self.norm(x, token_lens = token_lens)
+
+class SimpleViT(nn.Module):
+    def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, channels = 3, dim_head = 64, specialize_qkv_depth = None):
+        super().__init__()
+        image_height, image_width = pair(image_size)
+        self.patch_size = patch_height, patch_width = pair(patch_size)
+
+        assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
+
+        patch_dim = channels * patch_height * patch_width
+
+        self.to_patch_embedding = nn.Sequential(
+            Rearrange("b c (h p1) (w p2) -> b (h w) (p1 p2 c)", p1 = patch_height, p2 = patch_width),
+            nn.LayerNorm(patch_dim),
+            nn.Linear(patch_dim, dim),
+            nn.LayerNorm(dim),
+        )
+
+        self.pos_embedding = posemb_sincos_2d(
+            h = image_height // patch_height,
+            w = image_width // patch_width,
+            dim = dim,
+        )
+
+        self.cls_token = nn.Parameter(torch.randn(dim) * 1e-2)
+
+        specialize_qkv_depth = default(specialize_qkv_depth, depth // 3) # author found just first third of transformer having specialized qkv projection for cls token is enough
+
+        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, specialize_qkv_depth)
+
+        self.pool = 'cls'
+        self.to_latent = nn.Identity()
+
+        self.linear_head = nn.Linear(dim, num_classes)
+
+    def forward(self, img):
+        device = img.device
+
+        x = self.to_patch_embedding(img)
+        x += self.pos_embedding.to(device, dtype=x.dtype)
+
+        b, n, _ = x.shape
+        cls_tokens = repeat(self.cls_token, 'd -> b 1 d', b = b)
+        x = cat((cls_tokens, x), dim = 1)
+
+        x = self.transformer(x, token_lens = (1, n))
+
+        x = x[:, 0]
+
+        x = self.to_latent(x)
+        return self.linear_head(x)
+
+if __name__ == '__main__':
+    v = SimpleViT(
+        image_size = 256,
+        patch_size = 32,
+        num_classes = 1000,
+        dim = 1024,
+        depth = 6,
+        heads = 16,
+        mlp_dim = 2048
+    )
+
+    img = torch.randn(1, 3, 256, 256)
+    out = v(img)
+
+    assert out.shape == (1, 1000)

vit_pytorch/simple_vit_with_value_residual.py

@@ -0,0 +1,159 @@
+import torch
+from torch import nn
+from torch.nn import Module, ModuleList
+
+from einops import rearrange
+from einops.layers.torch import Rearrange
+
+# helpers
+
+def exists(v):
+    return v is not None
+
+def default(v, d):
+    return v if exists(v) else d
+
+def pair(t):
+    return t if isinstance(t, tuple) else (t, t)
+
+def posemb_sincos_2d(h, w, dim, temperature: int = 10000, dtype = torch.float32):
+    y, x = torch.meshgrid(torch.arange(h), torch.arange(w), indexing="ij")
+    assert (dim % 4) == 0, "feature dimension must be multiple of 4 for sincos emb"
+    omega = torch.arange(dim // 4) / (dim // 4 - 1)
+    omega = 1.0 / (temperature ** omega)
+
+    y = y.flatten()[:, None] * omega[None, :]
+    x = x.flatten()[:, None] * omega[None, :]
+    pe = torch.cat((x.sin(), x.cos(), y.sin(), y.cos()), dim=1)
+    return pe.type(dtype)
+
+# classes
+
+def FeedForward(dim, hidden_dim):
+    return nn.Sequential(
+        nn.LayerNorm(dim),
+        nn.Linear(dim, hidden_dim),
+        nn.GELU(),
+        nn.Linear(hidden_dim, dim),
+    )
+
+class Attention(Module):
+    def __init__(self, dim, heads = 8, dim_head = 64, learned_value_residual_mix = False):
+        super().__init__()
+        inner_dim = dim_head *  heads
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+        self.norm = nn.LayerNorm(dim)
+
+        self.attend = nn.Softmax(dim = -1)
+
+        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
+        self.to_out = nn.Linear(inner_dim, dim, bias = False)
+
+        self.to_residual_mix = nn.Sequential(
+            nn.Linear(dim, heads),
+            nn.Sigmoid(),
+            Rearrange('b n h -> b h n 1')
+        ) if learned_value_residual_mix else (lambda _: 0.5)
+
+    def forward(self, x, value_residual = None):
+        x = self.norm(x)
+
+        qkv = self.to_qkv(x).chunk(3, dim = -1)
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
+
+        if exists(value_residual):
+            mix = self.to_residual_mix(x)
+            v = v * mix + value_residual * (1. - mix)
+
+        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+
+        attn = self.attend(dots)
+
+        out = torch.matmul(attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+
+        return self.to_out(out), v
+
+class Transformer(Module):
+    def __init__(self, dim, depth, heads, dim_head, mlp_dim):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.layers = ModuleList([])
+        for i in range(depth):
+            is_first = i == 0
+            self.layers.append(ModuleList([
+                Attention(dim, heads = heads, dim_head = dim_head, learned_value_residual_mix = not is_first),
+                FeedForward(dim, mlp_dim)
+            ]))
+    def forward(self, x):
+        value_residual = None
+
+        for attn, ff in self.layers:
+
+            attn_out, values = attn(x, value_residual = value_residual)
+            value_residual = default(value_residual, values)
+
+            x = attn_out + x
+            x = ff(x) + x
+
+        return self.norm(x)
+
+class SimpleViT(Module):
+    def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, channels = 3, dim_head = 64):
+        super().__init__()
+        image_height, image_width = pair(image_size)
+        self.patch_size = patch_height, patch_width = pair(patch_size)
+
+        assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
+
+        patch_dim = channels * patch_height * patch_width
+
+        self.to_patch_embedding = nn.Sequential(
+            Rearrange("b c (h p1) (w p2) -> b (h w) (p1 p2 c)", p1 = patch_height, p2 = patch_width),
+            nn.LayerNorm(patch_dim),
+            nn.Linear(patch_dim, dim),
+            nn.LayerNorm(dim),
+        )
+
+        self.pos_embedding = posemb_sincos_2d(
+            h = image_height // patch_height,
+            w = image_width // patch_width,
+            dim = dim,
+        )
+
+        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim)
+
+        self.pool = "mean"
+        self.to_latent = nn.Identity()
+
+        self.linear_head = nn.Linear(dim, num_classes)
+
+    def forward(self, img):
+        device = img.device
+
+        x = self.to_patch_embedding(img)
+        x += self.pos_embedding.to(device, dtype=x.dtype)
+
+        x = self.transformer(x)
+        x = x.mean(dim = 1)
+
+        x = self.to_latent(x)
+        return self.linear_head(x)
+
+# quick test
+
+if __name__ == '__main__':
+    v = SimpleViT(
+        num_classes = 1000,
+        image_size = 256,
+        patch_size = 8,
+        dim = 1024,
+        depth = 6,
+        heads = 8,
+        mlp_dim = 2048,
+    )
+
+    images = torch.randn(2, 3, 256, 256)
+
+    logits = v(images)

vit_pytorch/t2t.py

@@ -61,10 +61,7 @@ class T2TViT(nn.Module):
         self.pool = pool
         self.to_latent = nn.Identity()
 
-        self.mlp_head = nn.Sequential(
-            nn.LayerNorm(dim),
-            nn.Linear(dim, num_classes)
-        )
+        self.mlp_head = nn.Linear(dim, num_classes)
 
     def forward(self, img):
         x = self.to_patch_embedding(img)

vit_pytorch/twins_svt.py

@@ -42,20 +42,11 @@ class LayerNorm(nn.Module):
         mean = torch.mean(x, dim = 1, keepdim = True)
         return (x - mean) / (var + self.eps).sqrt() * self.g + self.b
 
-class PreNorm(nn.Module):
-    def __init__(self, dim, fn):
-        super().__init__()
-        self.norm = LayerNorm(dim)
-        self.fn = fn
-
-    def forward(self, x, **kwargs):
-        x = self.norm(x)
-        return self.fn(x, **kwargs)
-
 class FeedForward(nn.Module):
     def __init__(self, dim, mult = 4, dropout = 0.):
         super().__init__()
         self.net = nn.Sequential(
+            LayerNorm(dim),
             nn.Conv2d(dim, dim * mult, 1),
             nn.GELU(),
             nn.Dropout(dropout),
@@ -99,6 +90,7 @@ class LocalAttention(nn.Module):
         self.heads = heads
         self.scale = dim_head ** -0.5
 
+        self.norm = LayerNorm(dim)
         self.to_q = nn.Conv2d(dim, inner_dim, 1, bias = False)
         self.to_kv = nn.Conv2d(dim, inner_dim * 2, 1, bias = False)
 
@@ -108,6 +100,8 @@ class LocalAttention(nn.Module):
         )
 
     def forward(self, fmap):
+        fmap = self.norm(fmap)
+
         shape, p = fmap.shape, self.patch_size
         b, n, x, y, h = *shape, self.heads
         x, y = map(lambda t: t // p, (x, y))
@@ -132,6 +126,8 @@ class GlobalAttention(nn.Module):
         self.heads = heads
         self.scale = dim_head ** -0.5
 
+        self.norm = LayerNorm(dim)
+
         self.to_q = nn.Conv2d(dim, inner_dim, 1, bias = False)
         self.to_kv = nn.Conv2d(dim, inner_dim * 2, k, stride = k, bias = False)
 
@@ -143,6 +139,8 @@ class GlobalAttention(nn.Module):
         )
 
     def forward(self, x):
+        x = self.norm(x)
+
         shape = x.shape
         b, n, _, y, h = *shape, self.heads
         q, k, v = (self.to_q(x), *self.to_kv(x).chunk(2, dim = 1))
@@ -164,10 +162,10 @@ class Transformer(nn.Module):
         self.layers = nn.ModuleList([])
         for _ in range(depth):
             self.layers.append(nn.ModuleList([
-                Residual(PreNorm(dim, LocalAttention(dim, heads = heads, dim_head = dim_head, dropout = dropout, patch_size = local_patch_size))) if has_local else nn.Identity(),
-                Residual(PreNorm(dim, FeedForward(dim, mlp_mult, dropout = dropout))) if has_local else nn.Identity(),
-                Residual(PreNorm(dim, GlobalAttention(dim, heads = heads, dim_head = dim_head, dropout = dropout, k = global_k))),
-                Residual(PreNorm(dim, FeedForward(dim, mlp_mult, dropout = dropout)))
+                Residual(LocalAttention(dim, heads = heads, dim_head = dim_head, dropout = dropout, patch_size = local_patch_size)) if has_local else nn.Identity(),
+                Residual(FeedForward(dim, mlp_mult, dropout = dropout)) if has_local else nn.Identity(),
+                Residual(GlobalAttention(dim, heads = heads, dim_head = dim_head, dropout = dropout, k = global_k)),
+                Residual(FeedForward(dim, mlp_mult, dropout = dropout))
             ]))
     def forward(self, x):
         for local_attn, ff1, global_attn, ff2 in self.layers:

vit_pytorch/vaat.py

@@ -0,0 +1,817 @@
+# vision-audio-action transformer - vaat
+
+from __future__ import annotations
+from contextlib import nullcontext
+
+import torch
+import torch.nn.functional as F
+from torch import nn, cat, stack, arange, tensor
+from torch.nn import Module, ModuleList
+
+from torchaudio.transforms import Spectrogram
+
+from einops import rearrange, repeat, reduce, pack, unpack
+from einops.layers.torch import Rearrange
+
+# helpers
+
+def exists(v):
+    return v is not None
+
+def default(v, d):
+    return v if exists(v) else d
+
+def pair(t):
+    return t if isinstance(t, tuple) else (t, t)
+
+# 2d sinusoidal positional embedding
+# simple vit paper shows it is good enough compared to learned
+
+def posemb_sincos_2d(
+    patches,
+    temperature = 10000,
+    dtype = torch.float32
+):
+    _, h, w, dim, device, dtype = *patches.shape, patches.device, patches.dtype
+
+    y, x = torch.meshgrid(arange(h, device = device), torch.arange(w, device = device), indexing = 'ij')
+    assert (dim % 4) == 0, 'feature dimension must be multiple of 4 for sincos emb'
+
+    omega = arange(dim // 4, device = device) / (dim // 4 - 1)
+    omega = temperature ** -omega
+
+    y = y.flatten()[:, None] * omega[None, :]
+    x = x.flatten()[:, None] * omega[None, :]
+
+    pe = cat((x.sin(), x.cos(), y.sin(), y.cos()), dim = 1)
+    pe = pe.type(dtype)
+
+    return rearrange(pe, '(h w) d -> h w d', h = h, w = w)
+
+# classes
+
+class FiLM(Module):
+    def __init__(
+        self,
+        dim,
+    ):
+        super().__init__()
+        proj = nn.Linear(dim, dim * 2)
+
+        self.to_gamma_beta = nn.Sequential(
+            proj,
+            Rearrange('b (two d) -> two b 1 d', two = 2)
+        )
+
+        nn.init.zeros_(proj.weight)
+        nn.init.zeros_(proj.bias)
+
+    def forward(self, tokens, cond):
+        gamma, beta = self.to_gamma_beta(cond)
+
+        return tokens * gamma + beta
+
+class FeedForward(Module):
+    def __init__(
+        self,
+        dim,
+        hidden_dim,
+        dropout = 0.
+    ):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Linear(dim, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, dim),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+class Attention(Module):
+    def __init__(
+        self,
+        dim,
+        heads = 8,
+        dim_head = 64,
+        dropout = 0.,
+        dim_context = None,
+        cross_attend = False
+    ):
+        super().__init__()
+        dim_context = default(dim_context, dim)
+        inner_dim = dim_head *  heads
+        project_out = not (heads == 1 and dim_head == dim)
+
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+
+        self.norm = nn.LayerNorm(dim)
+
+        self.cross_attend = cross_attend
+        self.context_norm = nn.LayerNorm(dim_context) if cross_attend else None
+
+        self.attend = nn.Softmax(dim = -1)
+        self.dropout = nn.Dropout(dropout)
+
+        self.to_q = nn.Linear(dim, inner_dim, bias = False)
+        self.to_kv = nn.Linear(dim_context, inner_dim * 2, bias = False)
+
+        self.to_out_gates = nn.Sequential(
+            nn.Linear(dim, heads),
+            Rearrange('b ... h -> b h ... 1'),
+            nn.Sigmoid()
+        )
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, dim),
+            nn.Dropout(dropout)
+        ) if project_out else nn.Identity()
+
+    def forward(self, x, context = None):
+
+        assert not (self.cross_attend ^ exists(context)), 'context must be passed in if cross attending, or vice versa'
+
+        x = self.norm(x)
+
+        # handle norming of context for cross attention
+
+        kv_input = x
+
+        if self.cross_attend:
+            context = self.context_norm(context)
+            kv_input = context
+
+        # project for queries, keys, values
+
+        qkv = (self.to_q(x), *self.to_kv(kv_input).chunk(2, dim = -1))
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
+
+        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+
+        attn = self.attend(dots)
+        attn = self.dropout(attn)
+
+        out = torch.matmul(attn, v)
+
+        out = out * self.to_out_gates(x)
+
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)
+
+class Transformer(Module):
+    def __init__(
+        self,
+        dim,
+        depth,
+        heads,
+        dim_head,
+        mlp_dim,
+        dropout = 0.
+    ):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.layers = ModuleList([])
+
+        for _ in range(depth):
+            self.layers.append(ModuleList([
+                Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout),
+                FeedForward(dim, mlp_dim, dropout = dropout)
+            ]))
+
+    def forward(
+        self,
+        x,
+        return_hiddens = False
+    ):
+
+        hiddens = []
+
+        for attn, ff in self.layers:
+            hiddens.append(x)
+
+            x = attn(x) + x
+            x = ff(x) + x
+
+        x = self.norm(x)
+
+        if not return_hiddens:
+            return x
+
+        return x, hiddens
+
+class AST(Module):
+    # audio spectrogram transformer https://arxiv.org/abs/2104.01778
+
+    def __init__(
+        self,
+        dim,
+        depth,
+        mlp_dim,
+        num_classes = None,
+        patch_size = 16,
+        dim_head = 64,
+        heads = 8,
+        dropout = 0.,
+        accept_spec = False,
+        accept_spec_time_first = True,
+        spec_n_fft = 128,
+        spec_power = 2,
+        spec_win_length = 24,
+        spec_hop_length = None,
+        spec_pad = 0,
+        spec_center = True,
+        spec_pad_mode = 'reflect',
+        num_register_tokens = 4
+    ):
+        super().__init__()
+        self.dim = dim
+        self.depth = depth
+
+        self.patch_size = patch_height, patch_width = pair(patch_size)
+        patch_input_dim = patch_height * patch_width
+
+        self.patch_size = (patch_height, patch_width)
+
+        self.to_patch_tokens = nn.Sequential(
+            Rearrange('b (h p1) (w p2) -> b h w (p1 p2)', p1 = self.patch_size[0], p2 = self.patch_size[1]),
+            nn.LayerNorm(patch_input_dim),
+            nn.Linear(patch_input_dim, dim),
+            nn.LayerNorm(dim)
+        )
+
+        self.accept_spec = accept_spec
+        self.accept_spec_time_first = accept_spec_time_first
+
+        self.spec = Spectrogram(
+            n_fft = spec_n_fft,
+            power = spec_power,
+            win_length = spec_win_length,
+            hop_length = spec_hop_length,
+            pad = spec_pad,
+            center = spec_center,
+            pad_mode = spec_pad_mode
+        )
+
+        self.transformer = Transformer(
+            dim = dim,
+            depth = depth,
+            dim_head = dim_head,
+            heads = heads,
+            mlp_dim = mlp_dim,
+            dropout = dropout,
+        )
+
+        self.final_norm = nn.LayerNorm(dim)
+
+        self.mlp_head = nn.Linear(dim, num_classes) if exists(num_classes) else nn.Identity()
+
+        self.register_tokens = nn.Parameter(torch.randn(num_register_tokens, dim) * 1e-2)
+
+    def forward(
+        self,
+        raw_audio_or_spec, # (b t) | (b f t)
+        return_hiddens = False
+    ):
+        batch, device = raw_audio_or_spec.shape[0], raw_audio_or_spec.device
+
+        assert (self.accept_spec and raw_audio_or_spec.ndim == 3) or (not self.accept_spec and raw_audio_or_spec.ndim == 2)
+
+        if self.accept_spec:
+            spec = rearrange(raw_audio_or_spec, 'b t f -> b f t')
+        else:
+            spec = self.spec(raw_audio_or_spec)
+
+        # automatically crop if audio does not yield a 2d spectrogram that is divisible by patch sizes
+
+        height, width = spec.shape[-2:]
+        patch_height, patch_width = self.patch_size
+
+        rounded_height = height // patch_height * patch_height
+        rounded_width = width // patch_width * patch_width
+
+        spec = spec[..., :rounded_height, :rounded_width]
+
+        # to patches
+
+        tokens = self.to_patch_tokens(spec)
+
+        # get number of patches along height and width
+
+        _, num_patch_height, num_patch_width, _ = tokens.shape
+
+        # 2d sinusoidal positional embedding
+
+        tokens = tokens + posemb_sincos_2d(tokens)
+
+        tokens = rearrange(tokens, 'b ... c -> b (...) c')
+
+        # register tokens
+
+        register_tokens = repeat(self.register_tokens, 'n d -> b n d', b = batch)
+
+        tokens, packed_shape = pack((register_tokens, tokens), 'b * d')
+
+        # attention
+
+        attended, hiddens = self.transformer(tokens, return_hiddens = True)
+
+        # final global average and norm (most recent papers show this is superior to CLS token)
+
+        normed = self.final_norm(attended)
+
+        if return_hiddens:
+            return normed, stack(hiddens)
+
+        register_tokens, normed = unpack(normed, packed_shape, 'b * d')
+
+        pooled = reduce(normed, 'b n d -> b d', 'mean')
+
+        maybe_logits = self.mlp_head(pooled)
+
+        return maybe_logits
+
+class ViT(Module):
+    def __init__(
+        self,
+        *,
+        image_size,
+        patch_size,
+        num_classes,
+        dim,
+        depth,
+        heads,
+        mlp_dim,
+        pool = 'cls',
+        channels = 3,
+        dim_head = 64,
+        dropout = 0.,
+        emb_dropout = 0.,
+        num_register_tokens = 0
+    ):
+        super().__init__()
+        self.dim = dim
+        self.depth = depth
+
+        image_height, image_width = pair(image_size)
+        self.patch_size = patch_height, patch_width = pair(patch_size)
+
+        assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
+
+        num_patches = (image_height // patch_height) * (image_width // patch_width)
+        patch_dim = channels * patch_height * patch_width
+        assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'
+
+        self.to_patch_embedding = nn.Sequential(
+            Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_height, p2 = patch_width),
+            nn.LayerNorm(patch_dim),
+            nn.Linear(patch_dim, dim),
+            nn.LayerNorm(dim),
+        )
+
+        self.pos_embedding = nn.Parameter(torch.randn(num_patches, dim))
+        self.cls_token = nn.Parameter(torch.randn(dim))
+        self.dropout = nn.Dropout(emb_dropout)
+
+        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)
+
+        self.pool = pool
+        self.to_latent = nn.Identity()
+
+        self.mlp_head = nn.Linear(dim, num_classes)
+
+        self.register_tokens = nn.Parameter(torch.randn(num_register_tokens, dim) * 1e-2)
+
+    def forward(self, img, return_hiddens = False):
+        x = self.to_patch_embedding(img)
+        b, n, _ = x.shape
+
+        x += self.pos_embedding[:n]
+
+        cls_tokens = repeat(self.cls_token, 'd -> b d', b = b)
+        register_tokens = repeat(self.register_tokens, 'n d -> b n d', b = b)
+
+        x, packed_shape = pack((register_tokens, cls_tokens, x), 'b * d')
+
+        x = self.dropout(x)
+
+        x, hiddens = self.transformer(x, return_hiddens = True)
+
+        # return the representation trajectory
+
+        if return_hiddens:
+            return x, stack(hiddens)
+
+        register_tokens, cls_tokens, x = unpack(x, packed_shape, 'b * d')
+
+        x = x.mean(dim = 1) if self.pool == 'mean' else cls_tokens
+
+        x = self.to_latent(x)
+
+        return self.mlp_head(x)
+
+# proposed VAT
+
+# https://openreview.net/forum?id=TalHOvvLZu
+# simple way to get SOTA on Libero dataset (beating fine-tuned pi-zero)
+
+class VAAT(Module):
+    def __init__(
+        self,
+        vit: ViT | dict,
+        ast: AST | dict,
+        *,
+        dim,
+        depth,
+        heads,
+        dim_head,
+        dim_action,
+        mlp_dim,
+        num_image_views = None,
+        num_audio_views = None,
+        num_tasks = None,
+        dim_extra_token = None,
+        num_register_tokens = 4,
+        action_chunk_len = 7,
+        time_seq_len = 1,
+        dropout = 0.,
+        add_self_attn = True,  # in the paper, they didn't have any ways for the action token to exchange information with the extra token, so we'll just add it as an option
+        self_attn_heads = 4,
+        self_attn_dim_head = 32,
+        ast_layer_indices: tuple[int, ...] | None = None,
+        vit_layer_indices: tuple[int, ...] | None = None,
+        num_advantage_bins = 0
+    ):
+        super().__init__()
+
+        # vit
+
+        if isinstance(vit, dict):
+            vit = ViT(**vit)
+
+        self.vit = vit
+
+        vit_dim = vit.dim
+
+        assert vit.depth == depth or exists(vit_layer_indices), f'if the VAAT depth is not equal to the ViT depth, you must pass in the indices from the ViT to be layered to the VAAT in order from bottom to top'
+
+        vit_layer_indices = default(vit_layer_indices, tuple(range(depth)))
+
+        assert len(vit_layer_indices) == depth, f'number of vit layer indices {len(vit_layer_indices)} does not much the VAT depth {depth}'
+
+        self.register_buffer('vit_layer_indices', tensor(vit_layer_indices), persistent = False)
+
+        # ast
+
+        if isinstance(ast, dict):
+            ast = AST(**ast)
+
+        self.ast = ast
+
+        ast_dim = ast.dim
+
+        self.ast_accept_spec = ast.accept_spec
+
+        assert ast.depth == depth or exists(ast_layer_indices), f'if the VAAT depth is not equal to the AST depth, you must pass in the indices from the AST to be layered to the VAAT in order from bottom to top'
+
+        ast_layer_indices = default(ast_layer_indices, tuple(range(depth)))
+
+        assert len(ast_layer_indices) == depth, f'number of ast layer indices {len(ast_layer_indices)} does not much the VAAT depth {depth}'
+
+        self.register_buffer('ast_layer_indices', tensor(ast_layer_indices), persistent = False)
+
+        # handle maybe multiple frames
+
+        is_video = time_seq_len > 1
+
+        self.is_video = is_video
+        self.time_seq_len = time_seq_len
+        self.time_pos_emb = nn.Parameter(torch.randn(time_seq_len, vit_dim) * 1e-2) if is_video else None
+
+        # maybe view embeddings
+
+        self.image_view_emb = nn.Parameter(torch.randn(num_image_views, vit_dim) * 1e-2) if exists(num_image_views) and num_image_views > 1 else None
+
+        self.audio_view_emb = nn.Parameter(torch.randn(num_audio_views, ast_dim) * 1e-2) if exists(num_audio_views) and num_audio_views > 1 else None
+
+        # handle maybe task conditioning
+
+        self.has_tasks = exists(num_tasks)
+
+        if self.has_tasks:
+            self.task_emb = nn.Parameter(torch.randn(num_tasks, dim) * 1e-2)
+
+        # register tokens from Darcet et al.
+
+        self.register_tokens = nn.Parameter(torch.randn(num_register_tokens, dim) * 1e-2)
+
+        # to action tokens
+
+        self.action_pos_emb = nn.Parameter(torch.randn(action_chunk_len, dim) * 1e-2)
+
+        # handle maybe advantage conditioning
+
+        self.has_advantages = num_advantage_bins > 0
+        self.num_advantage_bins = num_advantage_bins
+
+        if self.has_advantages:
+            self.advantage_emb = nn.Embedding(num_advantage_bins + 1, dim)
+
+        self.layers = ModuleList([])
+
+        for _ in range(depth):
+            maybe_film = FiLM(dim = dim) if self.has_tasks else None
+            maybe_self_attn = Attention(dim = dim, heads = self_attn_heads, dim_head = self_attn_dim_head, dropout = dropout) if add_self_attn else None
+
+            self.layers.append(ModuleList([
+                maybe_film,
+                maybe_self_attn,
+                Attention(dim = dim, dim_context = vit_dim, heads = heads, dim_head = dim_head, dropout = dropout, cross_attend = True),
+                Attention(dim = dim, dim_context = ast_dim, heads = heads, dim_head = dim_head, dropout = dropout, cross_attend = True),
+                FeedForward(dim = dim, hidden_dim = mlp_dim, dropout = dropout)
+            ]))
+
+        self.final_norm = nn.LayerNorm(dim)
+        self.to_pred_action = nn.Linear(dim, dim_action, bias = False)
+
+        # handle the extra token
+
+        self.accept_extra_token = exists(dim_extra_token)
+
+        if exists(dim_extra_token):
+            self.to_extra_token = nn.Linear(dim_extra_token, dim)
+
+    def forward(
+        self,
+        video_or_image,   # (b v? c t? h w)      - batch, views [wrist + third person or more], channels, maybe time, height, width
+        audio_or_spec,    # (b v? t) | (b v?f t) - batch, audio len | batch, spec freq, time
+        *,
+        extra = None,     # (b d)                - batch, dim extra
+        tasks = None,     # (b)
+        advantages = None,# (b)
+        actions = None,   # (b k d)              - batch, action chunk length, action dimension
+        return_hiddens = False,
+        freeze_vit = False,
+        freeze_ast = False
+    ):
+        batch, device = video_or_image.shape[0], video_or_image.device
+        return_loss = exists(actions)
+
+        # handle some various input dimensions
+
+        if video_or_image.ndim == 4:
+            video_or_image = rearrange(video_or_image, 'b 1 c h w')
+
+        assert (
+            (video_or_image.ndim == 5 and not self.is_video) or
+            (video_or_image.ndim == 6 and self.is_video)
+        )
+
+        if video_or_image.ndim == 5:
+            video_or_image = rearrange(video_or_image, 'b v c h w -> b v c 1 h w')
+
+        assert video_or_image.shape[3] == self.time_seq_len
+
+        # audio shapes - adding view if impliciy to be 1
+
+        if audio_or_spec.ndim == 2 and not self.ast_accept_spec:
+            audio_or_spec = rearrange(audio_or_spec, 'b t -> b 1 t')
+
+        elif audio_or_spec.ndim == 3 and self.ast_accept_spec:
+            audio_or_spec = rearrange(audio_or_spec, 'b f t -> b 1 f t')
+
+        # to images
+
+        images = rearrange(video_or_image, 'b v c t h w -> b v t c h w')
+
+        images, image_packed_shape = pack([images], '* c h w')
+
+        # to audio
+
+        if self.ast_accept_spec:
+            audio_or_spec, audio_packed_shape = pack([audio_or_spec], '* f t')
+        else:
+            audio_or_spec, audio_packed_shape = pack([audio_or_spec], '* t')
+
+        # get representation trajectory from vit
+
+        vit_forward_context = torch.no_grad if freeze_vit else nullcontext
+
+        with vit_forward_context():
+            embed, hiddens = self.vit(images, return_hiddens = True)
+
+        hiddens = cat((hiddens, embed[None, ...]))
+
+        # extract the hiddens needed for the action cross attention
+
+        hiddens = hiddens[self.vit_layer_indices]
+
+        # unpack temporarily for embedding
+
+        hiddens, = unpack(hiddens, image_packed_shape, 'l * n d') # l for layers
+
+        # maybe add time embeddings
+
+        if self.is_video:
+            time_pos_emb = rearrange(self.time_pos_emb, 't d -> t 1 d')
+            hiddens = hiddens + time_pos_emb
+
+        # maybe view embeddings
+
+        if exists(self.image_view_emb):
+            assert self.image_view_emb.shape[0] == hiddens.shape[2]
+
+            image_view_emb = rearrange(self.image_view_emb, 'v d -> v 1 1 d')
+            hiddens = hiddens + image_view_emb
+
+        # get representation trajectory from ast
+
+        ast_forward_context = torch.no_grad if freeze_ast else nullcontext
+
+        with ast_forward_context():
+            audio_embed, audio_hiddens = self.ast(audio_or_spec, return_hiddens = True)
+
+        audio_hiddens = cat((audio_hiddens, audio_embed[None, ...]))
+
+        # extract the hiddens needed for the action cross attention
+
+        audio_hiddens = audio_hiddens[self.ast_layer_indices]
+
+        # unpack audio temporarily for embedding
+
+        audio_hiddens, = unpack(audio_hiddens, audio_packed_shape, 'l * n d') # l for layers
+
+        # maybe audio view embeddings
+
+        if exists(self.audio_view_emb):
+            assert self.audio_view_emb.shape[0] == audio_hiddens.shape[2]
+
+            audio_view_emb = rearrange(self.audio_view_emb, 'v d -> v 1 1 d')
+            audio_hiddens = audio_hiddens + audio_view_emb
+
+        # maybe tasks
+
+        if exists(tasks):
+            assert self.has_tasks, f'`num_tasks` must be set on `VAT` for task conditioning'
+
+            task_emb = self.task_emb[tasks]
+
+        # cross from actions to representation trajectory
+
+        image_context = rearrange(hiddens, 'l b v t n d -> l b (v t n) d')
+
+        audio_context = rearrange(audio_hiddens, 'l b v n d -> l b (v n) d')
+
+        # main action tokens
+
+        action_tokens = repeat(self.action_pos_emb, 'n d -> b n d', b = batch)
+
+        # maybe advantage tokens
+
+        empty_token = action_tokens[:, 0:0]
+
+        maybe_advantage_embed = empty_token
+
+        if self.has_advantages and exists(advantages):
+            if isinstance(advantages, int):
+                advantages = torch.full((batch,), advantages, device = device, dtype = torch.long)
+
+            maybe_advantage_embed = self.advantage_emb(advantages + 1)
+
+        # register tokens
+
+        register_tokens = empty_token
+
+        if exists(self.register_tokens):
+            register_tokens = repeat(self.register_tokens, 'n d -> b n d', b = batch)
+
+        # extra
+
+        maybe_extra_embed = empty_token
+
+        has_extra = exists(extra)
+        if has_extra:
+            assert self.accept_extra_token
+
+            maybe_extra_embed = self.to_extra_token(extra)
+
+        # pack all tokens for attention
+
+        tokens, ps = pack((register_tokens, maybe_advantage_embed, action_tokens, maybe_extra_embed), 'b * d')
+
+        # transformer
+
+        hiddens = [tokens]
+
+        for (maybe_film, maybe_self_attn, image_cross_attn, audio_cross_attn, ff), image_layer_context, audio_layer_context in zip(self.layers, image_context, audio_context):
+
+            if exists(maybe_film) and exists(tasks):
+                tokens = maybe_film(tokens, task_emb)
+
+            tokens = image_cross_attn(tokens, image_layer_context) + tokens
+
+            tokens = audio_cross_attn(tokens, audio_layer_context) + tokens
+
+            if exists(maybe_self_attn):
+                tokens = maybe_self_attn(tokens) + tokens
+
+            tokens = ff(tokens) + tokens
+
+            hiddens.append(tokens)
+
+        # unpack register, advantage, action, and extra tokens
+
+        maybe_register_embed, maybe_advantage_embed, action_tokens, maybe_extra_embed = unpack(tokens, ps, 'b * d')
+
+        # norm and prediction
+
+        action_tokens = self.final_norm(action_tokens)
+
+        pred_action = self.to_pred_action(action_tokens)
+
+        if not return_loss:
+            if not return_hiddens:
+                return pred_action
+
+            return pred_action, stack(hiddens)
+
+        assert pred_action.shape[1] == actions.shape[1]
+
+        # they found l1 loss suffices
+
+        return F.l1_loss(pred_action, actions)
+
+# quick test
+
+if __name__ == '__main__':
+
+    vit = ViT(
+        image_size = 256,
+        patch_size = 32,
+        num_classes = 1000,
+        dim = 384,
+        heads = 8,
+        depth = 4,
+        mlp_dim = 384 * 4
+    )
+
+    ast = AST(
+        dim = 384,
+        depth = 4,
+        heads = 8,
+        num_classes = 1000,
+        patch_size = 16,
+        mlp_dim = 384 * 4
+    )
+
+    for num_adv_bins in (0, 2, 10):
+        vaat = VAAT(
+            vit,
+            ast,
+            dim = 512,
+            depth = 9,
+            heads = 8,
+            dim_head = 64,
+            mlp_dim = 2048,
+            dim_action = 20,
+            action_chunk_len = 7,
+            time_seq_len = 4,
+            num_image_views = 2,
+            num_audio_views = 2,
+            num_tasks = 4,
+            num_advantage_bins = num_adv_bins,
+            add_self_attn = True,
+            dim_extra_token = 33,               # extra token with some variable dimension
+            vit_layer_indices = (               # extending on the paper, allow for any order of hiddens, and also allow for depth index (which equates to the final embedding output from the vit)
+                0, 0, 1, 1, 2, 2, 3, 3, 4
+            ),
+            ast_layer_indices = (
+                1, 1, 1, 2, 2, 2, 3, 3, 3
+            )
+        )
+
+        images = torch.randn(2, 2, 3, 4, 256, 256) # (2 views with 4 frames)
+        audio = torch.randn(2, 2, 14_100 * 5)
+
+        tasks = torch.randint(0, 4, (2,))
+        extra = torch.randn(2, 33)                 # extra internal state
+
+        # advantage conditioning
+
+        advantages = None
+        if num_adv_bins > 0:
+            advantages = torch.randint(-1, num_adv_bins, (2,))
+
+        actions = torch.randn(2, 7, 20)         # actions for learning
+
+        loss = vaat(images, audio, actions = actions, advantages = advantages, tasks = tasks, extra = extra, freeze_vit = True)
+        loss.backward()
+
+        # after much training
+
+        pred_actions, hiddens = vaat(images, audio, advantages = advantages, tasks = tasks, extra = extra, return_hiddens = True)
+
+        assert pred_actions.shape == (2, 7, 20)

vit_pytorch/vat.py

@@ -0,0 +1,567 @@
+from __future__ import annotations
+from contextlib import nullcontext
+
+import torch
+import torch.nn.functional as F
+from torch import nn, cat, stack, tensor
+from torch.nn import Module, ModuleList
+
+from einops import rearrange, repeat, pack, unpack
+from einops.layers.torch import Rearrange
+
+# helpers
+
+def exists(v):
+    return v is not None
+
+def default(v, d):
+    return v if exists(v) else d
+
+def pair(t):
+    return t if isinstance(t, tuple) else (t, t)
+
+# classes
+
+class FiLM(Module):
+    def __init__(
+        self,
+        dim,
+    ):
+        super().__init__()
+        proj = nn.Linear(dim, dim * 2)
+
+        self.to_gamma_beta = nn.Sequential(
+            proj,
+            Rearrange('b (two d) -> two b 1 d', two = 2)
+        )
+
+        nn.init.zeros_(proj.weight)
+        nn.init.zeros_(proj.bias)
+
+    def forward(self, tokens, cond):
+        gamma, beta = self.to_gamma_beta(cond)
+
+        return tokens * gamma + beta
+
+class FeedForward(Module):
+    def __init__(
+        self,
+        dim,
+        hidden_dim,
+        dropout = 0.
+    ):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Linear(dim, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, dim),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+class Attention(Module):
+    def __init__(
+        self,
+        dim,
+        dim_context = None,
+        heads = 8,
+        dim_head = 64,
+        dropout = 0.,
+        cross_attend = False
+    ):
+        super().__init__()
+        dim_context = default(dim_context, dim)
+        inner_dim = dim_head *  heads
+        project_out = not (heads == 1 and dim_head == dim)
+
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+
+        self.norm = nn.LayerNorm(dim)
+
+        self.cross_attend = cross_attend
+        self.context_norm = nn.LayerNorm(dim_context) if cross_attend else None
+
+        self.attend = nn.Softmax(dim = -1)
+        self.dropout = nn.Dropout(dropout)
+
+        self.to_q = nn.Linear(dim, inner_dim, bias = False)
+        self.to_kv = nn.Linear(dim_context, inner_dim * 2, bias = False)
+
+        self.to_out_gates = nn.Sequential(
+            nn.Linear(dim, heads),
+            Rearrange('b ... h -> b h ... 1'),
+            nn.Sigmoid()
+        )
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, dim),
+            nn.Dropout(dropout)
+        ) if project_out else nn.Identity()
+
+    def forward(self, x, context = None):
+
+        assert not (self.cross_attend ^ exists(context)), 'context must be passed in if cross attending, or vice versa'
+
+        x = self.norm(x)
+
+        # handle norming of context for cross attention
+
+        kv_input = x
+
+        if self.cross_attend:
+            context = self.context_norm(context)
+            kv_input = context
+
+        # project for queries, keys, values
+
+        qkv = (self.to_q(x), *self.to_kv(kv_input).chunk(2, dim = -1))
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
+
+        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+
+        attn = self.attend(dots)
+        attn = self.dropout(attn)
+
+        out = torch.matmul(attn, v)
+        out = out * self.to_out_gates(x) # https://arxiv.org/abs/2505.06708
+
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)
+
+class Transformer(Module):
+    def __init__(
+        self,
+        dim,
+        depth,
+        heads,
+        dim_head,
+        mlp_dim,
+        dropout = 0.
+    ):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.layers = ModuleList([])
+
+        for _ in range(depth):
+            self.layers.append(ModuleList([
+                Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout),
+                FeedForward(dim, mlp_dim, dropout = dropout)
+            ]))
+
+    def forward(
+        self,
+        x,
+        return_hiddens = False
+    ):
+
+        hiddens = []
+
+        for attn, ff in self.layers:
+            hiddens.append(x)
+
+            x = attn(x) + x
+            x = ff(x) + x
+
+        x = self.norm(x)
+
+        if not return_hiddens:
+            return x
+
+        return x, hiddens
+
+class ViT(Module):
+    def __init__(
+        self,
+        *,
+        image_size,
+        patch_size,
+        num_classes,
+        dim,
+        depth,
+        heads,
+        mlp_dim,
+        pool = 'cls',
+        channels = 3,
+        dim_head = 64,
+        dropout = 0.,
+        emb_dropout = 0.,
+        num_register_tokens = 0
+    ):
+        super().__init__()
+        self.dim = dim
+        self.depth = depth
+
+        image_height, image_width = pair(image_size)
+        self.patch_size = patch_height, patch_width = pair(patch_size)
+
+        assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
+
+        num_patches = (image_height // patch_height) * (image_width // patch_width)
+        patch_dim = channels * patch_height * patch_width
+        assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'
+
+        self.to_patch_embedding = nn.Sequential(
+            Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_height, p2 = patch_width),
+            nn.LayerNorm(patch_dim),
+            nn.Linear(patch_dim, dim),
+            nn.LayerNorm(dim),
+        )
+
+        self.pos_embedding = nn.Parameter(torch.randn(num_patches, dim))
+        self.cls_token = nn.Parameter(torch.randn(dim))
+        self.dropout = nn.Dropout(emb_dropout)
+
+        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)
+
+        self.pool = pool
+        self.to_latent = nn.Identity()
+
+        self.mlp_head = nn.Linear(dim, num_classes)
+
+        self.register_tokens = nn.Parameter(torch.randn(num_register_tokens, dim) * 1e-2)
+
+    def forward(self, img, return_hiddens = False):
+        x = self.to_patch_embedding(img)
+        b, n, _ = x.shape
+
+        x += self.pos_embedding[:n]
+
+        cls_tokens = repeat(self.cls_token, 'd -> b d', b = b)
+        register_tokens = repeat(self.register_tokens, 'n d -> b n d', b = b)
+
+        x, packed_shape = pack((register_tokens, cls_tokens, x), 'b * d')
+
+        x = self.dropout(x)
+
+        x, hiddens = self.transformer(x, return_hiddens = True)
+
+        # return the representation trajectory
+
+        if return_hiddens:
+            return x, stack(hiddens)
+
+        register_tokens, cls_tokens, x = unpack(x, packed_shape, 'b * d')
+
+        x = x.mean(dim = 1) if self.pool == 'mean' else cls_tokens
+
+        x = self.to_latent(x)
+        return self.mlp_head(x)
+
+# proposed VAT
+
+# https://openreview.net/forum?id=TalHOvvLZu
+# simple way to get SOTA on Libero dataset (beating fine-tuned pi-zero)
+
+class VAT(Module):
+    def __init__(
+        self,
+        vit: ViT | dict,
+        *,
+        dim,
+        depth,
+        heads,
+        dim_head,
+        dim_action,
+        mlp_dim,
+        num_views = None,
+        num_tasks = None,
+        dim_extra_token = None,
+        num_register_tokens = 4,
+        action_chunk_len = 7,
+        time_seq_len = 1,
+        dropout = 0.,
+        add_self_attn = True,  # in the paper, they didn't have any ways for the action token to exchange information with the extra token, so we'll just add it as an option
+        self_attn_heads = 4,
+        self_attn_dim_head = 32,
+        vit_layer_indices: tuple[int, ...] | None = None,
+        num_advantage_bins = 0
+    ):
+        super().__init__()
+
+        if isinstance(vit, dict):
+            vit = ViT(**vit)
+
+        self.vit = vit
+
+        vit_dim = vit.dim
+
+        assert vit.depth == depth or exists(vit_layer_indices), f'if the VAT depth is not equal to the ViT depth, you must pass in the indices from the ViT to be layered to the VAT in order from bottom to top'
+
+        vit_layer_indices = default(vit_layer_indices, tuple(range(depth)))
+
+        assert len(vit_layer_indices) == depth, f'number of vit layer indices {len(vit_layer_indices)} does not much the VAT depth {depth}'
+
+        self.register_buffer('layer_indices', tensor(vit_layer_indices), persistent = False)
+
+        # handle maybe multiple frames
+
+        is_video = time_seq_len > 1
+
+        self.is_video = is_video
+        self.time_seq_len = time_seq_len
+        self.time_pos_emb = nn.Parameter(torch.randn(time_seq_len, vit_dim) * 1e-2) if is_video else None
+
+        # maybe view embeddings
+
+        self.view_emb = nn.Parameter(torch.randn(num_views, vit_dim) * 1e-2) if exists(num_views) and num_views > 1 else None
+
+        # handle maybe task conditioning
+
+        self.has_tasks = exists(num_tasks)
+
+        if self.has_tasks:
+            self.task_emb = nn.Parameter(torch.randn(num_tasks, dim) * 1e-2)
+
+        # register tokens from Darcet et al.
+
+        self.register_tokens = nn.Parameter(torch.randn(num_register_tokens, dim) * 1e-2)
+
+        # to action tokens
+
+        self.action_pos_emb = nn.Parameter(torch.randn(action_chunk_len, dim) * 1e-2)
+
+        # handle maybe advantage conditioning
+
+        self.has_advantages = num_advantage_bins > 0
+        self.num_advantage_bins = num_advantage_bins
+
+        if self.has_advantages:
+            self.advantage_emb = nn.Embedding(num_advantage_bins + 1, dim)
+
+        self.layers = ModuleList([])
+
+        for _ in range(depth):
+            maybe_film = FiLM(dim = dim) if self.has_tasks else None
+            maybe_self_attn = Attention(dim = dim, heads = self_attn_heads, dim_head = self_attn_dim_head, dropout = dropout) if add_self_attn else None
+
+            self.layers.append(ModuleList([
+                maybe_film,
+                maybe_self_attn,
+                Attention(dim = dim, dim_context = vit_dim, heads = heads, dim_head = dim_head, dropout = dropout, cross_attend = True),
+                FeedForward(dim = dim, hidden_dim = mlp_dim, dropout = dropout)
+            ]))
+
+        self.final_norm = nn.LayerNorm(dim)
+        self.to_pred_action = nn.Linear(dim, dim_action, bias = False)
+
+        # handle the extra token
+
+        self.accept_extra_token = exists(dim_extra_token)
+
+        if exists(dim_extra_token):
+            self.to_extra_token = nn.Linear(dim_extra_token, dim)
+
+    def forward(
+        self,
+        video_or_image,   # (b v? c t? h w) - batch, views [wrist + third person or more], channels, maybe time, height, width
+        *,
+        extra = None,     # (b d)           - batch, dim extra
+        tasks = None,     # (b)
+        advantages = None,# (b)
+        actions = None,   # (b k d)         - batch, action chunk length, action dimension
+        return_hiddens = False,
+        freeze_vit = False
+    ):
+        batch, device = video_or_image.shape[0], video_or_image.device
+        return_loss = exists(actions)
+
+        # handle some various input dimensions
+
+        if video_or_image.ndim == 4:
+            video_or_image = rearrange(video_or_image, 'b 1 c h w')
+
+        assert (
+            (video_or_image.ndim == 5 and not self.is_video) or
+            (video_or_image.ndim == 6 and self.is_video)
+        )
+
+        if video_or_image.ndim == 5:
+            video_or_image = rearrange(video_or_image, 'b v c h w -> b v c 1 h w')
+
+        assert video_or_image.shape[3] == self.time_seq_len
+
+        # to images
+
+        images = rearrange(video_or_image, 'b v c t h w -> b v t c h w')
+
+        images, packed_shape = pack([images], '* c h w')
+
+        # get representation trajectory from vit
+
+        vit_forward_context = torch.no_grad if freeze_vit else nullcontext
+
+        with vit_forward_context():
+            embed, hiddens = self.vit(images, return_hiddens = True)
+
+        hiddens = cat((hiddens, embed[None, ...]))
+
+        # extract the hiddens needed for the action cross attention
+
+        hiddens = hiddens[self.layer_indices]
+
+        # pack temporarily for embedding
+
+        hiddens, = unpack(hiddens, packed_shape, 'l * n d') # l for layers
+
+        # maybe add time embeddings
+
+        if self.is_video:
+            time_pos_emb = rearrange(self.time_pos_emb, 't d -> t 1 d')
+            hiddens = hiddens + time_pos_emb
+
+        # maybe view embeddings
+
+        if exists(self.view_emb):
+            assert self.view_emb.shape[0] == hiddens.shape[2]
+
+            view_emb = rearrange(self.view_emb, 'v d -> v 1 1 d')
+            hiddens = hiddens + view_emb
+
+        # maybe tasks
+
+        if exists(tasks):
+            assert self.has_tasks, f'`num_tasks` must be set on `VAT` for task conditioning'
+
+            task_emb = self.task_emb[tasks]
+
+        # cross from actions to representation trajectory
+
+        context = rearrange(hiddens, 'l b v t n d -> l b (v t n) d')
+
+        # main action tokens
+
+        action_tokens = repeat(self.action_pos_emb, 'n d -> b n d', b = batch)
+
+        # maybe advantage tokens
+
+        empty_token = action_tokens[:, 0:0]
+
+        maybe_advantage_embed = empty_token
+
+        if self.has_advantages and exists(advantages):
+            if isinstance(advantages, int):
+                advantages = torch.full((batch,), advantages, device = device, dtype = torch.long)
+
+            maybe_advantage_embed = self.advantage_emb(advantages + 1)
+
+        # register tokens
+
+        register_tokens = empty_token
+
+        if exists(self.register_tokens):
+            register_tokens = repeat(self.register_tokens, 'n d -> b n d', b = batch)
+
+        # extra
+
+        maybe_extra_embed = empty_token
+
+        has_extra = exists(extra)
+        if has_extra:
+            assert self.accept_extra_token
+
+            maybe_extra_embed = self.to_extra_token(extra)
+
+        # pack all tokens for attention
+
+        tokens, ps = pack((register_tokens, maybe_advantage_embed, action_tokens, maybe_extra_embed), 'b * d')
+
+        # transformer
+
+        hiddens = [tokens]
+
+        for (maybe_film, maybe_self_attn, cross_attn, ff), layer_context in zip(self.layers, context):
+
+            if exists(maybe_film) and exists(tasks):
+                tokens = maybe_film(tokens, task_emb)
+
+            tokens = cross_attn(tokens, layer_context) + tokens
+
+            if exists(maybe_self_attn):
+                tokens = maybe_self_attn(tokens) + tokens
+
+            tokens = ff(tokens) + tokens
+
+            hiddens.append(tokens)
+
+        # unpack register, advantage, action, and extra tokens
+
+        maybe_register_embed, maybe_advantage_embed, action_tokens, maybe_extra_embed = unpack(tokens, ps, 'b * d')
+
+        # norm and prediction
+
+        action_tokens = self.final_norm(action_tokens)
+
+        pred_action = self.to_pred_action(action_tokens)
+
+        if not return_loss:
+            if not return_hiddens:
+                return pred_action
+
+            return pred_action, stack(hiddens)
+
+        assert pred_action.shape[1] == actions.shape[1]
+
+        # they found l1 loss suffices
+
+        return F.l1_loss(pred_action, actions)
+
+# quick test
+
+if __name__ == '__main__':
+
+    vit = ViT(
+        image_size = 256,
+        patch_size = 32,
+        num_classes = 1000,
+        dim = 256,
+        heads = 8,
+        depth = 4,
+        mlp_dim = 1024
+    )
+
+    for num_adv_bins in (0, 2, 10):
+        vat = VAT(
+            vit,
+            dim = 512,
+            depth = 9,
+            heads = 8,
+            dim_head = 64,
+            mlp_dim = 2048,
+            dim_action = 20,
+            action_chunk_len = 7,
+            time_seq_len = 4,
+            num_views = 2,
+            num_tasks = 4,
+            num_advantage_bins = num_adv_bins,
+            add_self_attn = True,
+            dim_extra_token = 33,               # extra token with some variable dimension
+            vit_layer_indices = (               # extending on the paper, allow for any order of hiddens, and also allow for depth index (which equates to the final embedding output from the vit)
+                0, 0, 1, 1, 2, 2, 3, 3, 4
+            )
+        )
+
+        images = torch.randn(2, 2, 3, 4, 256, 256) # (2 views with 4 frames)
+        tasks = torch.randint(0, 4, (2,))
+        extra = torch.randn(2, 33)                 # extra internal state
+
+        # advantage conditioning
+
+        advantages = None
+        if num_adv_bins > 0:
+            advantages = torch.randint(-1, num_adv_bins, (2,))
+
+        actions = torch.randn(2, 7, 20)         # actions for learning
+
+        loss = vat(images, actions = actions, advantages = advantages, tasks = tasks, extra = extra, freeze_vit = True)
+        loss.backward()
+
+        # after much training
+
+        pred_actions, hiddens = vat(images, advantages = advantages, tasks = tasks, extra = extra, return_hiddens = True)
+
+        assert pred_actions.shape == (2, 7, 20)

vit_pytorch/vat_siglip.py

@@ -0,0 +1,521 @@
+from __future__ import annotations
+from contextlib import nullcontext
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+from torch import nn, cat, stack, tensor, einsum
+from torch.nn import Module, ModuleList
+
+from einops import rearrange, repeat, pack, unpack
+from einops.layers.torch import Rearrange
+
+# helpers
+
+def exists(v):
+    return v is not None
+
+def default(v, d):
+    return v if exists(v) else d
+
+def pair(t):
+    return t if isinstance(t, tuple) else (t, t)
+
+# attention
+
+class Attention(Module):
+    def __init__(
+        self,
+        dim,
+        dim_context = None,
+        heads = 8,
+        dim_head = 64,
+        dropout = 0.,
+        norm_eps = 1e-6,
+        gate_attn = False
+    ):
+        super().__init__()
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+        inner_dim = dim_head * heads
+
+        self.norm = nn.LayerNorm(dim, eps = norm_eps)
+
+        self.is_cross_attn = exists(dim_context)
+        dim_context = default(dim_context, dim)
+        self.norm_context = nn.LayerNorm(dim_context, eps = norm_eps) if self.is_cross_attn else None
+
+        self.to_q = nn.Linear(dim, inner_dim)
+        self.to_kv = nn.Linear(dim_context, inner_dim * 2)
+
+        self.to_out_gates = nn.Sequential(
+            nn.Linear(dim, heads),
+            Rearrange('b ... h -> b h ... 1'),
+            nn.Sigmoid()
+        ) if gate_attn else None
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, dim),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x, context = None):
+        x = self.norm(x)
+
+        if self.is_cross_attn:
+            assert exists(context)
+            context = self.norm_context(context)
+        else:
+            context = x
+
+        q = self.to_q(x)
+        k, v = self.to_kv(context).chunk(2, dim = -1)
+
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), (q, k, v))
+
+        sim = einsum('b h i d, b h j d -> b h i j', q, k) * self.scale
+        attn = sim.softmax(dim = -1)
+
+        out = einsum('b h i j, b h j d -> b h i d', attn, v)
+
+        if exists(self.to_out_gates):
+            out = out * self.to_out_gates(x) # https://arxiv.org/abs/2505.06708
+
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)
+
+def FeedForward(
+    dim,
+    dim_inner,
+    norm_eps = 1e-6
+):
+    return nn.Sequential(
+        nn.LayerNorm(dim, eps = norm_eps),
+        nn.Linear(dim, dim_inner),
+        nn.GELU(approximate = 'tanh'),
+        nn.Linear(dim_inner, dim)
+    )
+
+class SigLIP(Module):
+    def __init__(
+        self,
+        image_size = 224,
+        patch_size = 14,
+        dim = 1152,
+        depth = 27,
+        heads = 16,
+        mlp_dim = 4304,
+        norm_eps = 1e-6
+    ):
+        super().__init__()
+        self.dim = dim
+        self.depth = depth
+        num_patches = (image_size // patch_size) ** 2
+        dim_head = dim // heads
+
+        self.to_patch_embed = nn.Sequential(
+            Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_size, p2 = patch_size),
+            nn.Linear(patch_size * patch_size * 3, dim)
+        )
+
+        self.pos_embed = nn.Parameter(torch.randn(num_patches, dim))
+
+        self.layers = ModuleList([])
+        for _ in range(depth):
+            self.layers.append(ModuleList([
+                Attention(dim, heads = heads, dim_head = dim_head, norm_eps = norm_eps),
+                FeedForward(dim = dim, dim_inner = mlp_dim, norm_eps = norm_eps)
+            ]))
+
+        self.norm = nn.LayerNorm(dim, eps = norm_eps)
+
+    def forward(self, x, return_hiddens = False):
+        x = self.to_patch_embed(x)
+        num_patches = x.shape[1]
+
+        x = x + self.pos_embed[:num_patches]
+
+        hiddens = []
+
+        for attn, ff in self.layers:
+            hiddens.append(x)
+
+            x = attn(x) + x
+            x = ff(x) + x
+
+        out = self.norm(x)
+
+        if not return_hiddens:
+            return out
+
+        return out, stack(hiddens)
+
+class FiLM(Module):
+    def __init__(self, dim):
+        super().__init__()
+        proj = nn.Linear(dim, dim * 2)
+
+        self.to_gamma_beta = nn.Sequential(
+            proj,
+            Rearrange('b (two d) -> two b 1 d', two = 2)
+        )
+
+        nn.init.zeros_(proj.weight)
+        nn.init.zeros_(proj.bias)
+
+    def forward(self, tokens, cond):
+        gamma, beta = self.to_gamma_beta(cond)
+        return tokens * gamma + beta
+
+class SigLIPVAT(Module):
+    def __init__(
+        self,
+        *,
+        dim = 512,
+        depth = 27,
+        heads = 8,
+        dim_head = 64,
+        dim_action = 32,
+        mlp_dim = 2048,
+        num_views = 1,
+        num_tasks = None,
+        dim_extra_token = None,
+        num_register_tokens = 4,
+        action_chunk_len = 50,
+        time_seq_len = 1,
+        dropout = 0.,
+        add_self_attn = True,
+        self_attn_heads = 4,
+        self_attn_dim_head = 32,
+        vit_layer_indices: tuple[int, ...] | None = None,
+        num_advantage_bins = 0,
+        siglip_image_size = 224,
+        siglip_patch_size = 14,
+        siglip_dim = 1152,
+        siglip_depth = 27,
+        siglip_heads = 16,
+        siglip_mlp_dim = 4304,
+        siglip_norm_eps = 1e-6,
+    ):
+        super().__init__()
+
+        self.vit = SigLIP(
+            image_size = siglip_image_size,
+            patch_size = siglip_patch_size,
+            dim = siglip_dim,
+            depth = siglip_depth,
+            heads = siglip_heads,
+            mlp_dim = siglip_mlp_dim,
+            norm_eps = siglip_norm_eps
+        )
+
+        vit_dim = siglip_dim
+        self.vit_dim = vit_dim
+
+        vit_layer_indices = default(vit_layer_indices, tuple(range(depth)))
+        assert len(vit_layer_indices) == depth, f'number of vit layer indices {len(vit_layer_indices)} does not much the VAT depth {depth}'
+        self.register_buffer('layer_indices', tensor(vit_layer_indices), persistent = False)
+
+        # handle maybe multiple frames
+
+        is_video = time_seq_len > 1
+        self.is_video = is_video
+        self.time_seq_len = time_seq_len
+        self.time_pos_emb = nn.Parameter(torch.randn(time_seq_len, vit_dim) * 1e-2) if is_video else None
+
+        # maybe view embeddings
+
+        self.view_emb = nn.Parameter(torch.randn(num_views, vit_dim) * 1e-2) if exists(num_views) and num_views > 1 else None
+
+        # handle maybe task conditioning
+
+        self.has_tasks = exists(num_tasks)
+        if self.has_tasks:
+            self.task_emb = nn.Parameter(torch.randn(num_tasks, dim) * 1e-2)
+
+        # register tokens
+
+        self.register_tokens = nn.Parameter(torch.randn(num_register_tokens, dim) * 1e-2)
+
+        # to action tokens
+
+        self.action_pos_emb = nn.Parameter(torch.randn(action_chunk_len, dim) * 1e-2)
+
+        # handle maybe advantage conditioning
+
+        self.has_advantages = num_advantage_bins > 0
+        self.num_advantage_bins = num_advantage_bins
+
+        if self.has_advantages:
+            self.advantage_emb = nn.Embedding(num_advantage_bins + 1, dim)
+
+        self.layers = ModuleList([])
+        for _ in range(depth):
+            maybe_film = FiLM(dim = dim) if self.has_tasks else None
+            maybe_self_attn = Attention(dim = dim, heads = self_attn_heads, dim_head = self_attn_dim_head, dropout = dropout) if add_self_attn else None
+
+            self.layers.append(ModuleList([
+                maybe_film,
+                maybe_self_attn,
+                Attention(dim = dim, dim_context = vit_dim, heads = heads, dim_head = dim_head, dropout = dropout, gate_attn = True),
+                FeedForward(dim = dim, dim_inner = mlp_dim)
+            ]))
+
+        self.final_norm = nn.LayerNorm(dim)
+        self.to_pred_action = nn.Linear(dim, dim_action, bias = False)
+
+        # handle the extra token
+
+        self.accept_extra_token = exists(dim_extra_token)
+        if exists(dim_extra_token):
+            self.to_extra_token = nn.Linear(dim_extra_token, dim)
+
+    def load_siglip(
+        self,
+        repo_id = 'google/siglip-so400m-patch14-224',
+        folder = 'checkpoints/siglip'
+    ):
+        folder = Path(folder)
+        if not folder.exists():
+            from huggingface_hub import snapshot_download
+            snapshot_download(
+                repo_id = repo_id,
+                local_dir = folder,
+                allow_patterns = ['config.json', 'model.safetensors']
+            )
+
+        from safetensors import safe_open
+        weights_path = folder / 'model.safetensors'
+
+        # Auto-detect prefix based on keys
+        with safe_open(weights_path, framework = 'pt') as f:
+            keys = f.keys()
+
+            vi_p = ''
+            if any(k.startswith('paligemma_with_expert.paligemma.model.vision_tower.vision_model') for k in keys):
+                vi_p = 'paligemma_with_expert.paligemma.model.vision_tower.vision_model.'
+            elif any(k.startswith('vision_model') for k in keys):
+                vi_p = 'vision_model.'
+
+            pz_state = self.vit.state_dict()
+
+            def copy_weight_bias(pz_prefix, vi_prefix):
+                pz_state[f'{pz_prefix}.weight'].copy_(f.get_tensor(f'{vi_prefix}.weight'))
+                pz_state[f'{pz_prefix}.bias'].copy_(f.get_tensor(f'{vi_prefix}.bias'))
+
+            # patch embedding
+            patch_weight = rearrange(f.get_tensor(f'{vi_p}embeddings.patch_embedding.weight'), 'd c h w -> d (h w c)')
+            pz_state['to_patch_embed.1.weight'].copy_(patch_weight)
+            pz_state['to_patch_embed.1.bias'].copy_(f.get_tensor(f'{vi_p}embeddings.patch_embedding.bias'))
+
+            # position embedding
+            pz_state['pos_embed'].copy_(f.get_tensor(f'{vi_p}embeddings.position_embedding.weight'))
+
+            # transformer layers
+            for i in range(self.vit.depth):
+                v_pi = f'{vi_p}encoder.layers.{i}'
+                v_pz = f'layers.{i}'
+
+                # attention
+                copy_weight_bias(f'{v_pz}.0.norm', f'{v_pi}.layer_norm1')
+                copy_weight_bias(f'{v_pz}.0.to_q', f'{v_pi}.self_attn.q_proj')
+
+                vk, vv = [f.get_tensor(f'{v_pi}.self_attn.{x}_proj.weight') for x in ('k', 'v')]
+                bk, bv = [f.get_tensor(f'{v_pi}.self_attn.{x}_proj.bias') for x in ('k', 'v')]
+
+                pz_state[f'{v_pz}.0.to_kv.weight'].copy_(cat((vk, vv), dim = 0))
+                pz_state[f'{v_pz}.0.to_kv.bias'].copy_(cat((bk, bv), dim = 0))
+
+                copy_weight_bias(f'{v_pz}.0.to_out.0', f'{v_pi}.self_attn.out_proj')
+
+                # feedforward
+                copy_weight_bias(f'{v_pz}.1.0', f'{v_pi}.layer_norm2')
+                copy_weight_bias(f'{v_pz}.1.1', f'{v_pi}.mlp.fc1')
+                copy_weight_bias(f'{v_pz}.1.3', f'{v_pi}.mlp.fc2')
+
+            # post-layernorm
+            copy_weight_bias('norm', f'{vi_p}post_layernorm')
+
+            self.vit.load_state_dict(pz_state)
+
+        print(f'Successfully loaded SigLIP weights from {repo_id}')
+
+    def forward(
+        self,
+        video_or_image,   # (b v? c t? h w) - batch, views [wrist + third person or more], channels, maybe time, height, width
+        *,
+        extra = None,     # (b d)           - batch, dim extra
+        tasks = None,     # (b)
+        advantages = None,# (b)
+        actions = None,   # (b k d)         - batch, action chunk length, action dimension
+        return_hiddens = False,
+        freeze_vit = False
+    ):
+        batch, device = video_or_image.shape[0], video_or_image.device
+        return_loss = exists(actions)
+
+        # handle some various input dimensions
+
+        if video_or_image.ndim == 4:
+            video_or_image = rearrange(video_or_image, 'b 1 c h w')
+
+        if video_or_image.ndim == 5:
+            video_or_image = rearrange(video_or_image, 'b v c h w -> b v c 1 h w')
+
+        assert video_or_image.shape[3] == self.time_seq_len
+
+        # to images
+
+        images = rearrange(video_or_image, 'b v c t h w -> b v t c h w')
+        images, packed_shape = pack([images], '* c h w')
+
+        # get representation trajectory from vit
+
+        vit_forward_context = torch.no_grad if freeze_vit else nullcontext
+
+        with vit_forward_context():
+            embed, hiddens = self.vit(images, return_hiddens = True)
+
+        hiddens = cat((hiddens, embed[None, ...]))
+
+        # extract the hiddens needed for the action cross attention
+
+        hiddens = hiddens[self.layer_indices]
+
+        # pack temporarily for embedding
+
+        hiddens, = unpack(hiddens, packed_shape, 'l * n d') # l for layers
+
+        # maybe add time embeddings
+
+        if self.is_video:
+            time_pos_emb = rearrange(self.time_pos_emb, 't d -> t 1 d')
+            hiddens = hiddens + time_pos_emb
+
+        # maybe view embeddings
+
+        if exists(self.view_emb):
+            view_emb = rearrange(self.view_emb, 'v d -> v 1 1 d')
+            hiddens = hiddens + view_emb
+
+        # maybe tasks
+
+        if exists(tasks):
+            task_emb = self.task_emb[tasks]
+
+        # cross from actions to representation trajectory
+
+        context = rearrange(hiddens, 'l b v t n d -> l b (v t n) d')
+
+        # main action tokens
+
+        action_tokens = repeat(self.action_pos_emb, 'n d -> b n d', b = batch)
+
+        # maybe advantage tokens
+
+        empty_token = action_tokens[:, 0:0]
+
+        maybe_advantage_embed = empty_token
+
+        if self.has_advantages and exists(advantages):
+            if isinstance(advantages, int):
+                advantages = torch.full((batch,), advantages, device = device, dtype = torch.long)
+
+            maybe_advantage_embed = self.advantage_emb(advantages + 1)
+
+        # register tokens
+
+        register_tokens = empty_token
+
+        if exists(self.register_tokens):
+            register_tokens = repeat(self.register_tokens, 'n d -> b n d', b = batch)
+
+        # extra
+
+        maybe_extra_embed = empty_token
+
+        has_extra = exists(extra)
+        if has_extra:
+            maybe_extra_embed = self.to_extra_token(extra)
+
+        # pack all tokens for attention
+
+        tokens, ps = pack((register_tokens, maybe_advantage_embed, action_tokens, maybe_extra_embed), 'b * d')
+
+        # transformer
+
+        vat_hiddens = [tokens]
+
+        for (maybe_film, maybe_self_attn, cross_attn, ff), layer_context in zip(self.layers, context):
+
+            if exists(maybe_film) and exists(tasks):
+                tokens = maybe_film(tokens, task_emb)
+
+            tokens = cross_attn(tokens, layer_context) + tokens
+
+            if exists(maybe_self_attn):
+                tokens = maybe_self_attn(tokens) + tokens
+
+            tokens = ff(tokens) + tokens
+
+            vat_hiddens.append(tokens)
+
+        # unpack register, advantage, action, and extra tokens
+
+        maybe_register_embed, maybe_advantage_embed, action_tokens, maybe_extra_embed = unpack(tokens, ps, 'b * d')
+
+        # norm and prediction
+
+        action_tokens = self.final_norm(action_tokens)
+        pred_action = self.to_pred_action(action_tokens)
+
+        if not return_loss:
+            if not return_hiddens:
+                return pred_action
+
+            return pred_action, stack(vat_hiddens)
+
+        assert pred_action.shape[1] == actions.shape[1]
+
+        return F.l1_loss(pred_action, actions)
+
+# quick test
+
+if __name__ == '__main__':
+    for num_adv_bins in (0, 2, 10):
+        vat = SigLIPVAT(
+            num_tasks = 4,
+            dim_extra_token = 32,
+            time_seq_len = 2,
+            num_views = 2,
+            depth = 4,
+            num_advantage_bins = num_adv_bins,
+            vit_layer_indices = (               # extending on the paper, allow for any order of hiddens, and also allow for depth index (which equates to the final embedding output from the vit)
+                0, 1, 26, 27
+            )
+        )
+
+        vat.load_siglip() # load siglip weights from hf
+
+        # inputs
+
+        images = torch.randn(1, 2, 3, 2, 224, 224) # (b, v, c, t, h, w)
+        tasks = torch.randint(0, 4, (1,))
+        extra = torch.randn(1, 32)
+
+        # advantage conditioning
+
+        advantages = None
+        if num_adv_bins > 0:
+            advantages = torch.randint(-1, num_adv_bins, (1,))
+
+        actions = torch.randn(1, 50, 32) # actions for learning
+
+        loss = vat(images, actions = actions, advantages = advantages, tasks = tasks, extra = extra, freeze_vit = True)
+        loss.backward()
+
+        # after much training
+
+        pred_actions = vat(images, advantages = advantages, tasks = tasks, extra = extra)
+
+        assert pred_actions.shape == (1, 50, 32)

vit_pytorch/vit.py

@@ -1,5 +1,6 @@
 import torch
 from torch import nn
+from torch.nn import Module, ModuleList
 
 from einops import rearrange, repeat
 from einops.layers.torch import Rearrange
@@ -11,28 +12,22 @@ def pair(t):
 
 # classes
 
-class PreNorm(nn.Module):
-    def __init__(self, dim, fn):
-        super().__init__()
-        self.norm = nn.LayerNorm(dim)
-        self.fn = fn
-    def forward(self, x, **kwargs):
-        return self.fn(self.norm(x), **kwargs)
-
-class FeedForward(nn.Module):
+class FeedForward(Module):
     def __init__(self, dim, hidden_dim, dropout = 0.):
         super().__init__()
         self.net = nn.Sequential(
+            nn.LayerNorm(dim),
             nn.Linear(dim, hidden_dim),
             nn.GELU(),
             nn.Dropout(dropout),
             nn.Linear(hidden_dim, dim),
             nn.Dropout(dropout)
         )
+
     def forward(self, x):
         return self.net(x)
 
-class Attention(nn.Module):
+class Attention(Module):
     def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):
         super().__init__()
         inner_dim = dim_head *  heads
@@ -41,6 +36,8 @@ class Attention(nn.Module):
         self.heads = heads
         self.scale = dim_head ** -0.5
 
+        self.norm = nn.LayerNorm(dim)
+
         self.attend = nn.Softmax(dim = -1)
         self.dropout = nn.Dropout(dropout)
 
@@ -52,6 +49,8 @@ class Attention(nn.Module):
         ) if project_out else nn.Identity()
 
     def forward(self, x):
+        x = self.norm(x)
+
         qkv = self.to_qkv(x).chunk(3, dim = -1)
         q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
 
@@ -64,32 +63,38 @@ class Attention(nn.Module):
         out = rearrange(out, 'b h n d -> b n (h d)')
         return self.to_out(out)
 
-class Transformer(nn.Module):
+class Transformer(Module):
     def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0.):
         super().__init__()
-        self.layers = nn.ModuleList([])
+        self.norm = nn.LayerNorm(dim)
+        self.layers = ModuleList([])
+
         for _ in range(depth):
-            self.layers.append(nn.ModuleList([
-                PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout)),
-                PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout))
+            self.layers.append(ModuleList([
+                Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout),
+                FeedForward(dim, mlp_dim, dropout = dropout)
             ]))
+
     def forward(self, x):
         for attn, ff in self.layers:
             x = attn(x) + x
             x = ff(x) + x
-        return x
 
-class ViT(nn.Module):
+        return self.norm(x)
+
+class ViT(Module):
     def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, pool = 'cls', channels = 3, dim_head = 64, dropout = 0., emb_dropout = 0.):
         super().__init__()
         image_height, image_width = pair(image_size)
-        patch_height, patch_width = pair(patch_size)
+        self.patch_size = patch_height, patch_width = pair(patch_size)
 
         assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
 
         num_patches = (image_height // patch_height) * (image_width // patch_width)
         patch_dim = channels * patch_height * patch_width
+
         assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'
+        num_cls_tokens = 1 if pool == 'cls' else 0
 
         self.to_patch_embedding = nn.Sequential(
             Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_height, p2 = patch_width),
@@ -98,8 +103,9 @@ class ViT(nn.Module):
             nn.LayerNorm(dim),
         )
 
-        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
-        self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
+        self.cls_token = nn.Parameter(torch.randn(num_cls_tokens, dim))
+        self.pos_embedding = nn.Parameter(torch.randn(num_patches + num_cls_tokens, dim))
+
         self.dropout = nn.Dropout(emb_dropout)
 
         self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)
@@ -107,22 +113,25 @@ class ViT(nn.Module):
         self.pool = pool
         self.to_latent = nn.Identity()
 
-        self.mlp_head = nn.Sequential(
-            nn.LayerNorm(dim),
-            nn.Linear(dim, num_classes)
-        )
+        self.mlp_head = nn.Linear(dim, num_classes) if num_classes > 0 else None
 
     def forward(self, img):
+        batch = img.shape[0]
         x = self.to_patch_embedding(img)
-        b, n, _ = x.shape
 
-        cls_tokens = repeat(self.cls_token, '1 1 d -> b 1 d', b = b)
-        x = torch.cat((cls_tokens, x), dim=1)
-        x += self.pos_embedding[:, :(n + 1)]
+        cls_tokens = repeat(self.cls_token, '... d -> b ... d', b = batch)
+        x = torch.cat((cls_tokens, x), dim = 1)
+
+        seq = x.shape[1]
+
+        x = x + self.pos_embedding[:seq]
         x = self.dropout(x)
 
         x = self.transformer(x)
 
+        if self.mlp_head is None:
+            return x
+
         x = x.mean(dim = 1) if self.pool == 'mean' else x[:, 0]
 
         x = self.to_latent(x)

vit_pytorch/vit_1d.py

@@ -6,18 +6,11 @@ from einops.layers.torch import Rearrange
 
 # classes
 
-class PreNorm(nn.Module):
-    def __init__(self, dim, fn):
-        super().__init__()
-        self.norm = nn.LayerNorm(dim)
-        self.fn = fn
-    def forward(self, x, **kwargs):
-        return self.fn(self.norm(x), **kwargs)
-
 class FeedForward(nn.Module):
     def __init__(self, dim, hidden_dim, dropout = 0.):
         super().__init__()
         self.net = nn.Sequential(
+            nn.LayerNorm(dim),
             nn.Linear(dim, hidden_dim),
             nn.GELU(),
             nn.Dropout(dropout),
@@ -36,6 +29,7 @@ class Attention(nn.Module):
         self.heads = heads
         self.scale = dim_head ** -0.5
 
+        self.norm = nn.LayerNorm(dim)
         self.attend = nn.Softmax(dim = -1)
         self.dropout = nn.Dropout(dropout)
 
@@ -47,6 +41,7 @@ class Attention(nn.Module):
         ) if project_out else nn.Identity()
 
     def forward(self, x):
+        x = self.norm(x)
         qkv = self.to_qkv(x).chunk(3, dim = -1)
         q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
 
@@ -65,8 +60,8 @@ class Transformer(nn.Module):
         self.layers = nn.ModuleList([])
         for _ in range(depth):
             self.layers.append(nn.ModuleList([
-                PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout)),
-                PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout))
+                Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout),
+                FeedForward(dim, mlp_dim, dropout = dropout)
             ]))
     def forward(self, x):
         for attn, ff in self.layers:

vit_pytorch/vit_3d.py

@@ -11,18 +11,11 @@ def pair(t):
 
 # classes
 
-class PreNorm(nn.Module):
-    def __init__(self, dim, fn):
-        super().__init__()
-        self.norm = nn.LayerNorm(dim)
-        self.fn = fn
-    def forward(self, x, **kwargs):
-        return self.fn(self.norm(x), **kwargs)
-
 class FeedForward(nn.Module):
     def __init__(self, dim, hidden_dim, dropout = 0.):
         super().__init__()
         self.net = nn.Sequential(
+            nn.LayerNorm(dim),
             nn.Linear(dim, hidden_dim),
             nn.GELU(),
             nn.Dropout(dropout),
@@ -41,6 +34,7 @@ class Attention(nn.Module):
         self.heads = heads
         self.scale = dim_head ** -0.5
 
+        self.norm = nn.LayerNorm(dim)
         self.attend = nn.Softmax(dim = -1)
         self.dropout = nn.Dropout(dropout)
 
@@ -52,6 +46,7 @@ class Attention(nn.Module):
         ) if project_out else nn.Identity()
 
     def forward(self, x):
+        x = self.norm(x)
         qkv = self.to_qkv(x).chunk(3, dim = -1)
         q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
 
@@ -70,8 +65,8 @@ class Transformer(nn.Module):
         self.layers = nn.ModuleList([])
         for _ in range(depth):
             self.layers.append(nn.ModuleList([
-                PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout)),
-                PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout))
+                Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout),
+                FeedForward(dim, mlp_dim, dropout = dropout)
             ]))
     def forward(self, x):
         for attn, ff in self.layers:
@@ -83,7 +78,7 @@ class ViT(nn.Module):
     def __init__(self, *, image_size, image_patch_size, frames, frame_patch_size, num_classes, dim, depth, heads, mlp_dim, pool = 'cls', channels = 3, dim_head = 64, dropout = 0., emb_dropout = 0.):
         super().__init__()
         image_height, image_width = pair(image_size)
-        patch_height, patch_width = pair(image_patch_size)
+        self.patch_size = patch_height, patch_width = pair(image_patch_size)
 
         assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
         assert frames % frame_patch_size == 0, 'Frames must be divisible by frame patch size'
@@ -94,7 +89,7 @@ class ViT(nn.Module):
         assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'
 
         self.to_patch_embedding = nn.Sequential(
-            Rearrange('b c (f pf) (h p1) (w p2) -> b (f h w) (p1 p2 pf c)', p1 = patch_height, p2 = patch_width, pf = frame_patch_size),
+            Rearrange('b c (f pf) (h p1) (w p2) -> b (f h w) (pf p1 p2 c)', p1 = patch_height, p2 = patch_width, pf = frame_patch_size),
             nn.LayerNorm(patch_dim),
             nn.Linear(patch_dim, dim),
             nn.LayerNorm(dim),

vit_pytorch/vit_5.py

@@ -0,0 +1,223 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.nn import Module, ModuleList
+from math import sqrt, pi
+
+from einops import rearrange, repeat
+from einops.layers.torch import Rearrange
+from torch.amp import autocast
+
+def pair(t):
+    return t if isinstance(t, tuple) else (t, t)
+
+def exists(val):
+    return val is not None
+
+@autocast('cuda', enabled = False)
+def rotate_every_two(x):
+    x = rearrange(x, '... (d j) -> ... d j', j = 2)
+    x1, x2 = x.unbind(dim = -1)
+    x = torch.stack((-x2, x1), dim = -1)
+    return rearrange(x, '... d j -> ... (d j)')
+
+class RMSNorm(Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.scale = dim ** 0.5
+        self.g = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        return F.normalize(x, dim=-1) * self.scale * self.g
+
+class FeedForward(Module):
+    def __init__(self, dim, hidden_dim, dropout = 0.):
+        super().__init__()
+        self.net = nn.Sequential(
+            RMSNorm(dim),
+            nn.Linear(dim, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, dim),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+class AxialRotaryEmbedding(nn.Module):
+    def __init__(self, dim, num_registers=4, base=100000, reg_base=100):
+        super().__init__()
+        self.dim = dim
+        self.num_registers = num_registers
+
+        half_dim = dim // 2
+
+        patch_inv_freq = 1.0 / (base ** (torch.arange(0, half_dim, 2).float() / half_dim))
+        self.register_buffer('patch_inv_freq', patch_inv_freq)
+
+        reg_inv_freq = 1.0 / (reg_base ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer('reg_inv_freq', reg_inv_freq)
+
+    @autocast('cuda', enabled=False)
+    def forward(self, x):
+        device, dtype = x.device, x.dtype
+        num_patches = x.shape[-2] - self.num_registers
+        n = int(sqrt(num_patches))
+
+        seq = torch.arange(n, device=device, dtype=dtype)
+
+        freqs = torch.einsum('i, j -> i j', seq, self.patch_inv_freq.to(dtype))
+
+        freqs_x = repeat(freqs, 'i d -> i j d', j=n)
+        freqs_y = repeat(freqs, 'j d -> i j d', i=n)
+
+        patch_freqs = torch.cat((freqs_x, freqs_y), dim=-1)
+        patch_freqs = rearrange(patch_freqs, 'i j d -> (i j) d')
+
+        # registers
+        reg_pos = torch.arange(self.num_registers, device=device, dtype=dtype)
+        reg_freqs = torch.einsum('i, j -> i j', reg_pos, self.reg_inv_freq.to(dtype))
+
+        all_freqs = torch.cat((reg_freqs, patch_freqs), dim=0)
+
+        sin = all_freqs.sin()
+        cos = all_freqs.cos()
+
+        sin, cos = map(lambda t: repeat(t, 'n d -> () () n (d j)', j=2), (sin, cos))
+        return sin, cos
+
+class Attention(Module):
+    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):
+        super().__init__()
+        inner_dim = dim_head * heads
+        project_out = not (heads == 1 and dim_head == dim)
+
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+
+        self.norm = RMSNorm(dim)
+
+        self.q_norm = RMSNorm(dim_head)
+        self.k_norm = RMSNorm(dim_head)
+
+        self.attend = nn.Softmax(dim = -1)
+        self.dropout = nn.Dropout(dropout)
+
+        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, dim),
+            nn.Dropout(dropout)
+        ) if project_out else nn.Identity()
+
+    def forward(self, x, pos_emb=None):
+        x = self.norm(x)
+
+        qkv = self.to_qkv(x).chunk(3, dim = -1)
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
+
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+
+        if exists(pos_emb):
+            sin, cos = pos_emb
+            (q_cls, q), (k_cls, k) = map(lambda t: (t[:, :, :1], t[:, :, 1:]), (q, k))
+            q, k = map(lambda t: (t * cos) + (rotate_every_two(t) * sin), (q, k))
+            q = torch.cat((q_cls, q), dim = 2)
+            k = torch.cat((k_cls, k), dim = 2)
+
+        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+
+        attn = self.attend(dots)
+        attn = self.dropout(attn)
+
+        out = torch.matmul(attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)
+
+class Transformer(Module):
+    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0.):
+        super().__init__()
+        self.norm = RMSNorm(dim)
+        self.layers = ModuleList([])
+
+        self.ls1 = nn.ParameterList([])
+        self.ls2 = nn.ParameterList([])
+
+        for _ in range(depth):
+            self.layers.append(ModuleList([
+                Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout),
+                FeedForward(dim, mlp_dim, dropout = dropout)
+            ]))
+            self.ls1.append(nn.Parameter(torch.ones(dim) * 1e-4))
+            self.ls2.append(nn.Parameter(torch.ones(dim) * 1e-4))
+
+    def forward(self, x, pos_emb=None):
+        for (attn, ff), ls1, ls2 in zip(self.layers, self.ls1, self.ls2):
+            x = attn(x, pos_emb=pos_emb) * ls1 + x
+            x = ff(x) * ls2 + x
+
+        return self.norm(x)
+
+class ViT(Module):
+    def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, channels = 3, dim_head = 64, dropout = 0., emb_dropout = 0., num_registers = 4):
+        super().__init__()
+        image_height, image_width = pair(image_size)
+        self.patch_size = patch_height, patch_width = pair(patch_size)
+
+        assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
+
+        num_patches = (image_height // patch_height) * (image_width // patch_width)
+        patch_dim = channels * patch_height * patch_width
+
+        self.num_registers = num_registers
+
+        self.to_patch_embedding = nn.Sequential(
+            Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_height, p2 = patch_width),
+            nn.Linear(patch_dim, dim),
+        )
+
+        self.cls_token = nn.Parameter(torch.randn(1, dim))
+
+        if num_registers > 0:
+            self.register_tokens = nn.Parameter(torch.randn(num_registers, dim))
+
+        self.pos_embedding = nn.Parameter(torch.randn(num_patches + 1 + num_registers, dim))
+
+        self.dropout = nn.Dropout(emb_dropout)
+
+        self.rotary_emb = AxialRotaryEmbedding(dim=dim_head, num_registers=num_registers)
+
+        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)
+
+        self.to_latent = nn.Identity()
+
+        self.mlp_head = nn.Linear(dim, num_classes) if num_classes > 0 else None
+
+    def forward(self, img):
+        batch = img.shape[0]
+        x = self.to_patch_embedding(img)
+
+        cls_tokens = repeat(self.cls_token, '... d -> b ... d', b = batch)
+
+        if self.num_registers > 0:
+            register_tokens = repeat(self.register_tokens, '... d -> b ... d', b = batch)
+            x = torch.cat((cls_tokens, register_tokens, x), dim = 1)
+        else:
+            x = torch.cat((cls_tokens, x), dim = 1)
+
+        seq = x.shape[1]
+
+        x = x + self.pos_embedding[:seq]
+        x = self.dropout(x)
+
+        pos_emb = self.rotary_emb(x[:, 1:])
+        x = self.transformer(x, pos_emb=pos_emb)
+
+        if self.mlp_head is None:
+            return x
+
+        x = x[:, 0]
+        x = self.to_latent(x)
+        return self.mlp_head(x)

vit_pytorch/vit_detpool.py

@@ -0,0 +1,254 @@
+from __future__ import annotations
+
+# DetPool ViT - a vit that accepts an object mask and attends and pools only using that mask - table 1
+# Dantong Niu et al. - https://openreview.net/forum?id=NZDaMcpXZm
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.nn import Module, ModuleList
+
+from einops import rearrange, repeat, pack, unpack
+from einops.layers.torch import Rearrange, Reduce
+
+# helpers
+
+def exists(val):
+    return val is not None
+
+def pair(t):
+    return t if isinstance(t, tuple) else (t, t)
+
+def masked_mean(t, mask, dim = 1, eps = 1e-5):
+    if not exists(mask):
+        return t.mean(dim = dim)
+
+    mask = rearrange(mask.bool(), '... -> ... 1')
+    t = t.masked_fill(~mask, 0.)
+    return t.sum(dim = dim) / mask.sum(dim = dim).clamp(min = eps)
+
+# classes
+
+class FeedForward(Module):
+    def __init__(self, dim, hidden_dim, dropout = 0.):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Linear(dim, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, dim),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+class Attention(Module):
+    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):
+        super().__init__()
+        inner_dim = dim_head *  heads
+        project_out = not (heads == 1 and dim_head == dim)
+
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+
+        self.norm = nn.LayerNorm(dim)
+
+        self.attend = nn.Softmax(dim = -1)
+        self.dropout = nn.Dropout(dropout)
+
+        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, dim),
+            nn.Dropout(dropout)
+        ) if project_out else nn.Identity()
+
+    def forward(self, x, mask = None):
+        x = self.norm(x)
+
+        qkv = self.to_qkv(x).chunk(3, dim = -1)
+        q, k, v = (rearrange(t, 'b n (h d) -> b h n d', h = self.heads) for t in qkv)
+
+        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+
+        if exists(mask):
+            mask = rearrange(mask, 'b j -> b 1 1 j')
+            mask_value = -torch.finfo(dots.dtype).max
+            dots = dots.masked_fill(~mask, mask_value)
+
+        attn = self.attend(dots)
+        attn = self.dropout(attn)
+
+        out = torch.matmul(attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)
+
+class Transformer(Module):
+    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0.):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.layers = ModuleList([])
+
+        for _ in range(depth):
+            self.layers.append(ModuleList([
+                Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout),
+                FeedForward(dim, mlp_dim, dropout = dropout)
+            ]))
+
+    def forward(self, x, mask = None):
+        for attn, ff in self.layers:
+            x = attn(x, mask = mask) + x
+            x = ff(x) + x
+
+        return self.norm(x)
+
+class ViTDetPool(Module):
+    def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, use_cls_token = True, channels = 3, dim_head = 64, dropout = 0., emb_dropout = 0., mask_generator: Module | None = None):
+        super().__init__()
+        image_height, image_width = pair(image_size)
+        self.patch_size = patch_height, patch_width = pair(patch_size)
+
+        assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
+
+        num_patches = (image_height // patch_height) * (image_width // patch_width)
+        patch_dim = channels * patch_height * patch_width
+
+        self.patch_height = patch_height
+        self.patch_width = patch_width
+
+        self.downsample_mask = Reduce('b (h p1) (w p2) -> b (h w)', 'max', p1 = patch_height, p2 = patch_width)
+
+        self.to_patch_embedding = nn.Sequential(
+            Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_height, p2 = patch_width),
+            nn.LayerNorm(patch_dim),
+            nn.Linear(patch_dim, dim),
+            nn.LayerNorm(dim),
+        )
+
+        # maybe cls
+
+        self.use_cls_token = use_cls_token
+
+        if use_cls_token:
+            self.cls_token = nn.Parameter(torch.randn(dim) * 1e-2)
+
+        self.pos_embedding = nn.Parameter(torch.randn(num_patches, dim) * 1e-2)
+
+        self.dropout = nn.Dropout(emb_dropout)
+
+        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)
+
+        self.to_latent = nn.Identity()
+
+        self.mlp_head = nn.Linear(dim, num_classes) if num_classes > 0 else None
+
+        self.mask_generator = mask_generator
+
+    def forward(self, img, object_mask = None):
+
+        if not exists(object_mask) and exists(self.mask_generator):
+            with torch.no_grad():
+                self.mask_generator.eval()
+                object_mask = self.mask_generator(img)
+
+        has_cls = self.use_cls_token
+
+        batch, _, height, width = img.shape
+        tokens = self.to_patch_embedding(img)
+
+        seq = tokens.shape[1]
+        tokens = tokens + self.pos_embedding[:seq]
+
+        if has_cls:
+            cls_token = repeat(self.cls_token, 'd -> b d', b = batch)
+            tokens, packed_shape = pack((cls_token, tokens), 'b * d')
+
+        tokens = self.dropout(tokens)
+
+        # handle the attention mask, and for final pooling
+
+        mask = None
+
+        if exists(object_mask):
+            assert object_mask.ndim in {3, 2}
+
+            if object_mask.shape == (batch, height, width):
+                mask = self.downsample_mask(object_mask)
+            else:
+                mask = object_mask
+
+            mask = rearrange(mask, 'b ... -> b (...)')
+
+            assert mask.shape == (batch, seq)
+
+            if has_cls:
+                mask = F.pad(mask, (1, 0), value = True)
+
+        # attend with maybe mask
+
+        tokens = self.transformer(tokens, mask = mask)
+
+        if not exists(self.mlp_head):
+            return tokens
+
+        # splice out cls
+
+        if has_cls:
+            _, tokens = unpack(tokens, packed_shape, 'b * d')
+
+            if exists(mask):
+                mask = mask[..., 1:]
+
+        # pooling with the mask
+
+        pooled = masked_mean(tokens, mask, dim = 1)
+
+        pooled = self.to_latent(pooled)
+        return self.mlp_head(pooled)
+
+# quick test
+
+if __name__ == '__main__':
+    vit = ViTDetPool(
+        image_size = 256,
+        patch_size = 32,
+        num_classes = 1000,
+        dim = 1024,
+        depth = 6,
+        heads = 16,
+        mlp_dim = 2048,
+        dropout = 0.1,
+        emb_dropout = 0.1
+    )
+
+    img = torch.randn(1, 3, 256, 256)
+    object_mask = torch.randint(0, 2, (1, 256, 256)).bool()
+
+    preds = vit(img, object_mask = object_mask)
+    assert preds.shape == (1, 1000)
+
+    preds_no_mask = vit(img)
+    assert preds_no_mask.shape == (1, 1000)
+
+    # test with module included
+
+    class MockMasker(Module):
+        def forward(self, img):
+            batch, _, height, width = img.shape
+            return torch.ones(batch, height, width).bool()
+
+    vit = ViTDetPool(
+        image_size = 256,
+        patch_size = 32,
+        num_classes = 1000,
+        dim = 1024,
+        depth = 1,
+        heads = 16,
+        mlp_dim = 2048,
+        mask_generator = MockMasker()
+    )
+
+    preds = vit(img)
+    assert preds.shape == (1, 1000)

vit_pytorch/vit_for_small_dataset.py

@@ -13,18 +13,11 @@ def pair(t):
 
 # classes
 
-class PreNorm(nn.Module):
-    def __init__(self, dim, fn):
-        super().__init__()
-        self.norm = nn.LayerNorm(dim)
-        self.fn = fn
-    def forward(self, x, **kwargs):
-        return self.fn(self.norm(x), **kwargs)
-
 class FeedForward(nn.Module):
     def __init__(self, dim, hidden_dim, dropout = 0.):
         super().__init__()
         self.net = nn.Sequential(
+            nn.LayerNorm(dim),
             nn.Linear(dim, hidden_dim),
             nn.GELU(),
             nn.Dropout(dropout),
@@ -41,6 +34,7 @@ class LSA(nn.Module):
         self.heads = heads
         self.temperature = nn.Parameter(torch.log(torch.tensor(dim_head ** -0.5)))
 
+        self.norm = nn.LayerNorm(dim)
         self.attend = nn.Softmax(dim = -1)
         self.dropout = nn.Dropout(dropout)
 
@@ -52,6 +46,7 @@ class LSA(nn.Module):
         )
 
     def forward(self, x):
+        x = self.norm(x)
         qkv = self.to_qkv(x).chunk(3, dim = -1)
         q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
 
@@ -74,8 +69,8 @@ class Transformer(nn.Module):
         self.layers = nn.ModuleList([])
         for _ in range(depth):
             self.layers.append(nn.ModuleList([
-                PreNorm(dim, LSA(dim, heads = heads, dim_head = dim_head, dropout = dropout)),
-                PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout))
+                LSA(dim, heads = heads, dim_head = dim_head, dropout = dropout),
+                FeedForward(dim, mlp_dim, dropout = dropout)
             ]))
     def forward(self, x):
         for attn, ff in self.layers:
@@ -104,7 +99,7 @@ class ViT(nn.Module):
     def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, pool = 'cls', channels = 3, dim_head = 64, dropout = 0., emb_dropout = 0.):
         super().__init__()
         image_height, image_width = pair(image_size)
-        patch_height, patch_width = pair(patch_size)
+        self.patch_size = patch_height, patch_width = pair(patch_size)
 
         assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
 

vit_pytorch/vit_nd.py

@@ -0,0 +1,191 @@
+from __future__ import annotations
+
+import torch
+from torch import nn
+from torch.nn import Module
+
+from einops import rearrange, repeat
+from einops.layers.torch import Rearrange
+
+# helpers
+
+def join(arr, delimiter = ' '):
+    return delimiter.join(arr)
+
+def ensure_tuple(t, length):
+    if isinstance(t, (tuple, list)):
+        assert len(t) == length, f'Expected tuple of length {length}, got {len(t)}'
+        return tuple(t)
+    return (t,) * length
+
+# classes
+
+class FeedForward(Module):
+    def __init__(self, dim, hidden_dim, dropout = 0.):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Linear(dim, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, dim),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+class Attention(Module):
+    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):
+        super().__init__()
+        inner_dim = dim_head * heads
+        project_out = not (heads == 1 and dim_head == dim)
+
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+
+        self.norm = nn.LayerNorm(dim)
+        self.attend = nn.Softmax(dim = -1)
+        self.dropout = nn.Dropout(dropout)
+
+        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, dim),
+            nn.Dropout(dropout)
+        ) if project_out else nn.Identity()
+
+    def forward(self, x):
+        x = self.norm(x)
+        qkv = self.to_qkv(x).chunk(3, dim = -1)
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
+
+        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+
+        attn = self.attend(dots)
+        attn = self.dropout(attn)
+
+        out = torch.matmul(attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)
+
+class Transformer(Module):
+    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0.):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.layers = nn.ModuleList([])
+        for _ in range(depth):
+            self.layers.append(nn.ModuleList([
+                Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout),
+                FeedForward(dim, mlp_dim, dropout = dropout)
+            ]))
+
+    def forward(self, x):
+        for attn, ff in self.layers:
+            x = attn(x) + x
+            x = ff(x) + x
+        return self.norm(x)
+
+class ViTND(Module):
+    def __init__(
+        self,
+        *,
+        ndim: int,
+        input_shape: int | tuple[int, ...],
+        patch_size: int | tuple[int, ...],
+        num_classes: int,
+        dim: int,
+        depth: int,
+        heads: int,
+        mlp_dim: int,
+        pool: str = 'cls',
+        channels: int = 3,
+        dim_head: int = 64,
+        dropout: float = 0.,
+        emb_dropout: float = 0.
+    ):
+        super().__init__()
+
+        assert 1 <= ndim <= 7, 'ndim must be between 1 and 7'
+        assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'
+
+        self.ndim = ndim
+        self.pool = pool
+
+        input_shape = ensure_tuple(input_shape, ndim)
+        patch_size = ensure_tuple(patch_size, ndim)
+
+        for i, (inp_dim, patch_dim) in enumerate(zip(input_shape, patch_size)):
+            assert inp_dim % patch_dim == 0, f'Input dimension {i} ({inp_dim}) must be divisible by patch size ({patch_dim})'
+
+        num_patches_per_dim = [inp_dim // patch_dim for inp_dim, patch_dim in zip(input_shape, patch_size)]
+        num_patches = 1
+        for n in num_patches_per_dim:
+            num_patches *= n
+
+        patch_dim = channels
+        for p in patch_size:
+            patch_dim *= p
+
+        dim_names = 'fghijkl'[:ndim]
+
+        input_dims = [f'({d} p{i})' for i, d in enumerate(dim_names)]
+        patch_dims = [f'p{i}' for i in range(ndim)]
+
+        input_pattern = f'b c {join(input_dims)}'
+        output_pattern = f'b ({join(dim_names)}) ({join(patch_dims)} c)'
+        rearrange_str = f'{input_pattern} -> {output_pattern}'
+
+        rearrange_kwargs = {f'p{i}': p for i, p in enumerate(patch_size)}
+
+        self.to_patch_embedding = nn.Sequential(
+            Rearrange(rearrange_str, **rearrange_kwargs),
+            nn.Linear(patch_dim, dim),
+            nn.LayerNorm(dim),
+        )
+
+        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
+        self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
+        self.dropout = nn.Dropout(emb_dropout)
+
+        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)
+
+        self.to_latent = nn.Identity()
+        self.mlp_head = nn.Linear(dim, num_classes)
+
+    def forward(self, x):
+        x = self.to_patch_embedding(x)
+        b, n, _ = x.shape
+
+        cls_tokens = repeat(self.cls_token, '1 1 d -> b 1 d', b = b)
+        x = torch.cat((cls_tokens, x), dim = 1)
+        x += self.pos_embedding[:, :(n + 1)]
+        x = self.dropout(x)
+
+        x = self.transformer(x)
+
+        x = x[:, 1:].mean(dim = 1) if self.pool == 'mean' else x[:, 0]
+
+        x = self.to_latent(x)
+        return self.mlp_head(x)
+
+
+if __name__ == '__main__':
+
+    model = ViTND(
+        ndim = 4,
+        input_shape = (8, 16, 32, 64),
+        patch_size = (2, 4, 4, 8),
+        num_classes = 1000,
+        dim = 512,
+        depth = 6,
+        heads = 8,
+        mlp_dim = 2048,
+        channels = 3,
+        dropout = 0.1,
+        emb_dropout = 0.1
+    )
+
+    occupancy_time = torch.randn(2, 3, 8, 16, 32, 64)
+
+    logits = model(occupancy_time)

vit_pytorch/vit_nd_pope.py

@@ -0,0 +1,353 @@
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from torch import pi, nn, arange, cat, stack, Tensor
+from torch.nn import Module, ModuleList
+from torch.amp import autocast
+
+from einops import rearrange, repeat, reduce, pack, unpack
+from einops.layers.torch import Rearrange
+
+# helpers
+
+def exists(val):
+    return val is not None
+
+def l2norm(t):
+    return F.normalize(t, dim = -1, p = 2)
+
+def join(arr, delimiter = ' '):
+    return delimiter.join(arr)
+
+def ensure_tuple(t, length):
+    if isinstance(t, (tuple, list)):
+        assert len(t) == length, f'Expected tuple of length {length}, got {len(t)}'
+        return tuple(t)
+
+    return (t,) * length
+
+# golden gate rotary - Jerry Xiong, PhD student at UIUC
+# https://jerryxio.ng/posts/nd-rope/
+
+# but using polar version instead
+# Gopalakrishnan et al. https://arxiv.org/abs/2509.10534
+
+def _phi(m: int) -> float:
+    x = 2.0
+    for _ in range(10):
+        x = (1 + x) ** (1.0 / (m + 1.0))
+    return x
+
+def make_directions(n: int, d: int) -> Tensor:
+    g = _phi(d)
+    alpha = (1.0 / g) ** arange(1, d + 1, dtype = torch.float64)
+    i = arange(1, n + 1, dtype = torch.float64).unsqueeze(1)
+    z = torch.fmod(i * alpha, 1.0)
+    directions = torch.erfinv(2.0 * z - 1.0)
+    directions = l2norm(directions)
+    return directions.float()
+
+class GoldenGatePoPENd(Module):
+    def __init__(
+        self,
+        dim_pos: int,
+        heads: int,
+        dim_head: int,
+        min_freq: float = 1.0,
+        max_freq: float = 10000.0,
+        p_zero_freqs: float = 0.0, # proportion of frequencies set to 0
+        init_learned_bias_uniform = False
+    ):
+        super().__init__()
+        n_freqs = dim_head
+        n_zero_freqs = round(p_zero_freqs * n_freqs)
+
+        omega = cat((
+            torch.zeros(n_zero_freqs),
+            min_freq * (max_freq / min_freq) ** torch.linspace(0, 1, n_freqs - n_zero_freqs),
+        ))
+
+        directions = rearrange(
+            make_directions(heads * n_freqs, dim_pos),
+            '(h f) p -> h f p',
+            h = heads
+        )
+
+        omega_expanded = rearrange(omega, 'f -> f 1')
+        self.register_buffer('freqs', directions * omega_expanded)  # shape: (h, f, p)
+
+        self.learned_bias = nn.Parameter(torch.zeros(heads, dim_head))
+
+        if init_learned_bias_uniform:
+            self.learned_bias.uniform_(-2. * pi, 0.)
+
+    @autocast('cuda', enabled = False)
+    def forward(self, pos):
+
+        freqs = rearrange(self.freqs, 'h f p -> 1 h 1 f p')
+        positions = rearrange(pos.float(), 'b n p -> b 1 n 1 p')
+
+        # compute theta for each (batch, head, seq, freq)
+
+        theta = reduce(freqs * positions, 'b h n f p -> b h n f', 'sum')
+
+        bias = self.learned_bias.clamp(-2. * pi, 0.)
+        bias = rearrange(bias, 'h d -> h 1 d')
+
+        return theta, bias
+
+@autocast('cuda', enabled = False)
+def apply_polar_pos_emb(t, freqs):
+    orig_dtype = t.dtype
+
+    t = t.float()
+    t = F.softplus(t)
+
+    out = cat((t * freqs.cos(), t * freqs.sin()), dim = -1)
+
+    return out.type(orig_dtype)
+
+# classes
+
+class FeedForward(Module):
+    def __init__(self, dim, hidden_dim, dropout = 0.):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Linear(dim, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, dim),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+class Attention(Module):
+    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):
+        super().__init__()
+        inner_dim = dim_head * heads
+        project_out = not (heads == 1 and dim_head == dim)
+
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+
+        self.norm = nn.LayerNorm(dim)
+        self.attend = nn.Softmax(dim = -1)
+        self.dropout = nn.Dropout(dropout)
+
+        self.to_qk = nn.Linear(dim, inner_dim * 2, bias = False)
+        self.to_v = nn.Linear(dim, inner_dim, bias = False)
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, dim),
+            nn.Dropout(dropout)
+        ) if project_out else nn.Identity()
+
+    def forward(self, x, polar_pos_emb = None):
+        x = self.norm(x)
+        qkv = (*self.to_qk(x).chunk(2, dim = -1), self.to_v(x))
+
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
+
+        if exists(polar_pos_emb):
+            freqs, bias = polar_pos_emb
+            q = apply_polar_pos_emb(q, freqs)
+            k = apply_polar_pos_emb(k, freqs + bias)
+
+        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+
+        attn = self.attend(dots)
+        attn = self.dropout(attn)
+
+        out = torch.matmul(attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)
+
+class Transformer(Module):
+    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0., polar_emb = None):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+
+        self.polar_emb = polar_emb
+
+        self.layers = ModuleList([])
+
+        for _ in range(depth):
+            self.layers.append(ModuleList([
+                Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout),
+                FeedForward(dim, mlp_dim, dropout = dropout)
+            ]))
+
+    def forward(self, x, pos = None):
+
+        # pope embedding
+
+        polar_pos_emb = None
+        if exists(pos) and exists(self.polar_emb):
+            polar_pos_emb = self.polar_emb(pos)
+
+        # transformer layers
+
+        for attn, ff in self.layers:
+            x = attn(x, polar_pos_emb) + x
+            x = ff(x) + x
+
+        return self.norm(x)
+
+class ViTND(Module):
+    def __init__(
+        self,
+        *,
+        ndim: int,
+        input_shape: int | tuple[int, ...],
+        patch_size: int | tuple[int, ...],
+        num_classes: int,
+        dim: int,
+        depth: int,
+        heads: int,
+        mlp_dim: int,
+        channels: int = 3,
+        dim_head: int = 64,
+        dropout: float = 0.,
+        emb_dropout: float = 0.,
+        pope_min_freq: float = 1.0,
+        pope_max_freq: float = 10000.0,
+        pope_p_zero_freqs: float = 0.0,
+        pope_init_learned_bias_uniform = False
+    ):
+        super().__init__()
+
+        assert 1 <= ndim <= 7, 'ndim must be between 1 and 7'
+
+        self.ndim = ndim
+
+        input_shape = ensure_tuple(input_shape, ndim)
+        patch_size = ensure_tuple(patch_size, ndim)
+
+        for i, (inp_dim, patch_dim) in enumerate(zip(input_shape, patch_size)):
+            assert inp_dim % patch_dim == 0, f'Input dimension {i} ({inp_dim}) must be divisible by patch size ({patch_dim})'
+
+        num_patches_per_dim = [inp_dim // patch_dim for inp_dim, patch_dim in zip(input_shape, patch_size)]
+        num_patches = 1
+        for n in num_patches_per_dim:
+            num_patches *= n
+
+        patch_dim = channels
+        for p in patch_size:
+            patch_dim *= p
+
+        dim_names = 'fghijkl'[:ndim]
+
+        input_dims = [f'({d} p{i})' for i, d in enumerate(dim_names)]
+        patch_dims = [f'p{i}' for i in range(ndim)]
+
+        input_pattern = f'b c {join(input_dims)}'
+        output_pattern = f'b {join(dim_names)} ({join(patch_dims)} c)'
+        rearrange_str = f'{input_pattern} -> {output_pattern}'
+
+        rearrange_kwargs = {f'p{i}': p for i, p in enumerate(patch_size)}
+
+        self.to_patch_embedding = nn.Sequential(
+            Rearrange(rearrange_str, **rearrange_kwargs),
+            nn.Linear(patch_dim, dim),
+            nn.LayerNorm(dim),
+        )
+
+        self.dropout = nn.Dropout(emb_dropout)
+
+        # golden gate pope
+
+        self.polar_emb = GoldenGatePoPENd(
+            dim_pos = ndim,
+            heads = heads,
+            dim_head = dim_head,
+            min_freq = pope_min_freq,
+            max_freq = pope_max_freq,
+            p_zero_freqs = pope_p_zero_freqs,
+            init_learned_bias_uniform = pope_init_learned_bias_uniform
+        )
+
+        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout, polar_emb = self.polar_emb)
+
+        self.to_latent = nn.Identity()
+        self.mlp_head = nn.Linear(dim, num_classes)
+
+    def muon_parameters(self):
+        params = []
+
+        for m in self.modules():
+            if isinstance(m, Attention):
+                params.extend([
+                    m.to_v.weight,
+                    m.to_out[0].weight
+                ])
+            elif isinstance(m, FeedForward):
+                params.extend([
+                    m.net[1].weight,
+                    m.net[-2].weight
+                ])
+
+        return params
+
+    def forward(
+        self,
+        x,
+        return_embed = False
+    ):
+        x = self.to_patch_embedding(x) # (b, *spatial_dims, patch_dim)
+
+        batch, *spatial_dims, _, device = *x.shape, x.device
+
+        # Generate position coordinates
+
+        grids = [arange(d, device = device, dtype = torch.float32) for d in spatial_dims]
+        grid = torch.meshgrid(*grids, indexing = 'ij')
+        pos = stack(grid, dim = -1)  # (*spatial_dims, ndim)
+
+        # flatten spatial dimensions for attention with nd rotary
+
+        pos = repeat(pos, '... p -> b (...) p', b = batch)
+        x, packed_shape = pack([x], 'b * d')
+
+        x = self.dropout(x)
+
+        embed = self.transformer(x, pos)
+
+        # return the embed with reconstituted patch shape
+
+        if return_embed:
+            embed, = unpack(embed, packed_shape, 'b * d')
+            return embed
+
+        # pooling to logits
+
+        pooled = reduce(embed, 'b n d -> b d', 'mean')
+
+        pooled = self.to_latent(pooled)
+        return self.mlp_head(pooled)
+
+if __name__ == '__main__':
+
+    model = ViTND(
+        ndim = 5,
+        input_shape = (4, 8, 16, 32, 64),
+        patch_size = (2, 2, 4, 4, 8),
+        num_classes = 1000,
+        dim = 512,
+        depth = 6,
+        heads = 8,
+        mlp_dim = 2048,
+        channels = 3,
+        dropout = 0.1,
+        emb_dropout = 0.1
+    )
+
+    data = torch.randn(3, 3, 4, 8, 16, 32, 64)
+
+    logits = model(data)
+
+    embed = model(data, return_embed = True)
+    assert embed.shape == (3, 2, 4, 4, 8, 8, 512)

vit_pytorch/vit_nd_rotary.py

@@ -0,0 +1,325 @@
+from __future__ import annotations
+
+import torch
+from torch import nn, arange, cat, stack, Tensor
+from torch.nn import Module, ModuleList
+import torch.nn.functional as F
+
+from einops import rearrange, repeat, reduce, pack, unpack
+from einops.layers.torch import Rearrange
+
+# helpers
+
+def exists(val):
+    return val is not None
+
+def l2norm(t):
+    return F.normalize(t, dim = -1, p = 2)
+
+def join(arr, delimiter = ' '):
+    return delimiter.join(arr)
+
+def ensure_tuple(t, length):
+    if isinstance(t, (tuple, list)):
+        assert len(t) == length, f'Expected tuple of length {length}, got {len(t)}'
+        return tuple(t)
+    return (t,) * length
+
+# golden gate rotary - Jerry Xiong, PhD student at UIUC
+# https://jerryxio.ng/posts/nd-rope/
+
+def _phi(m: int) -> float:
+    x = 2.0
+    for _ in range(10):
+        x = (1 + x) ** (1.0 / (m + 1.0))
+    return x
+
+def make_directions(n: int, d: int) -> Tensor:
+    g = _phi(d)
+    alpha = (1.0 / g) ** arange(1, d + 1, dtype = torch.float64)
+    i = arange(1, n + 1, dtype = torch.float64).unsqueeze(1)
+    z = torch.fmod(i * alpha, 1.0)
+    directions = torch.erfinv(2.0 * z - 1.0)
+    directions = l2norm(directions)
+    return directions.float()
+
+class GoldenGateRoPENd(Module):
+    def __init__(
+        self,
+        dim_pos: int,
+        heads: int,
+        dim_head: int,
+        rope_min_freq: float = 1.0,
+        rope_max_freq: float = 10000.0,
+        rope_p_zero_freqs: float = 0.0, # proportion of frequencies set to 0
+    ):
+        super().__init__()
+        n_freqs = dim_head // 2
+        n_zero_freqs = round(rope_p_zero_freqs * n_freqs)
+
+        omega = cat((
+            torch.zeros(n_zero_freqs),
+            rope_min_freq * (rope_max_freq / rope_min_freq) ** torch.linspace(0, 1, n_freqs - n_zero_freqs),
+        ))
+
+        directions = rearrange(
+            make_directions(heads * n_freqs, dim_pos),
+            '(h f) p -> h f p',
+            h = heads
+        )
+
+        omega_expanded = rearrange(omega, 'f -> f 1')
+        self.register_buffer('freqs', directions * omega_expanded)  # shape: (h, f, p)
+
+    def forward(self, input: Tensor, pos: Tensor) -> Tensor:
+        # input shape: (b, h, n, d) where d = head_dim
+        # pos shape: (b, n, p) where p = pos_dim
+        # self.freqs shape: (h, f, p) where f = d // 2
+
+        x, y = input.float().chunk(2, dim = -1)  # both (b, h, n, f)
+
+        # Expand dimensions for broadcasting
+        freqs = rearrange(self.freqs, 'h f p -> 1 h 1 f p')
+        positions = rearrange(pos.float(), 'b n p -> b 1 n 1 p')
+
+        # Compute theta for each (batch, head, seq, freq)
+        theta = reduce(freqs * positions, 'b h n f p -> b h n f', 'sum')
+
+        cos_theta = torch.cos(theta)
+        sin_theta = torch.sin(theta)
+
+        # Apply rotation
+        x_out = x * cos_theta - y * sin_theta
+        y_out = x * sin_theta + y * cos_theta
+
+        output = cat((x_out, y_out), dim=-1)
+        return output.type_as(input)
+
+# classes
+
+class FeedForward(Module):
+    def __init__(self, dim, hidden_dim, dropout = 0.):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Linear(dim, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, dim),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+class Attention(Module):
+    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0., rotary_emb = None):
+        super().__init__()
+        inner_dim = dim_head * heads
+        project_out = not (heads == 1 and dim_head == dim)
+
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+        self.rotary_emb = rotary_emb
+
+        self.norm = nn.LayerNorm(dim)
+        self.attend = nn.Softmax(dim = -1)
+        self.dropout = nn.Dropout(dropout)
+
+        self.to_qk = nn.Linear(dim, inner_dim * 2, bias = False)
+        self.to_v = nn.Linear(dim, inner_dim, bias = False)
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, dim),
+            nn.Dropout(dropout)
+        ) if project_out else nn.Identity()
+
+    def forward(self, x, pos = None):
+        x = self.norm(x)
+        qkv = (*self.to_qk(x).chunk(2, dim = -1), self.to_v(x))
+
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
+
+        # Apply rotary embeddings if available
+        if exists(self.rotary_emb):
+            assert exists(pos)
+            q = self.rotary_emb(q, pos)
+            k = self.rotary_emb(k, pos)
+
+        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+
+        attn = self.attend(dots)
+        attn = self.dropout(attn)
+
+        out = torch.matmul(attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)
+
+class Transformer(Module):
+    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0., rotary_emb = None):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.layers = ModuleList([])
+        for _ in range(depth):
+            self.layers.append(ModuleList([
+                Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout, rotary_emb = rotary_emb),
+                FeedForward(dim, mlp_dim, dropout = dropout)
+            ]))
+
+    def forward(self, x, pos = None):
+        for attn, ff in self.layers:
+            x = attn(x, pos) + x
+            x = ff(x) + x
+        return self.norm(x)
+
+class ViTND(Module):
+    def __init__(
+        self,
+        *,
+        ndim: int,
+        input_shape: int | tuple[int, ...],
+        patch_size: int | tuple[int, ...],
+        num_classes: int,
+        dim: int,
+        depth: int,
+        heads: int,
+        mlp_dim: int,
+        channels: int = 3,
+        dim_head: int = 64,
+        dropout: float = 0.,
+        emb_dropout: float = 0.,
+        rope_min_freq: float = 1.0,
+        rope_max_freq: float = 10000.0,
+        rope_p_zero_freqs: float = 0.0
+    ):
+        super().__init__()
+
+        assert 1 <= ndim <= 7, 'ndim must be between 1 and 7'
+
+        self.ndim = ndim
+
+        input_shape = ensure_tuple(input_shape, ndim)
+        patch_size = ensure_tuple(patch_size, ndim)
+
+        for i, (inp_dim, patch_dim) in enumerate(zip(input_shape, patch_size)):
+            assert inp_dim % patch_dim == 0, f'Input dimension {i} ({inp_dim}) must be divisible by patch size ({patch_dim})'
+
+        num_patches_per_dim = [inp_dim // patch_dim for inp_dim, patch_dim in zip(input_shape, patch_size)]
+        num_patches = 1
+        for n in num_patches_per_dim:
+            num_patches *= n
+
+        patch_dim = channels
+        for p in patch_size:
+            patch_dim *= p
+
+        dim_names = 'fghijkl'[:ndim]
+
+        input_dims = [f'({d} p{i})' for i, d in enumerate(dim_names)]
+        patch_dims = [f'p{i}' for i in range(ndim)]
+
+        input_pattern = f'b c {join(input_dims)}'
+        output_pattern = f'b {join(dim_names)} ({join(patch_dims)} c)'
+        rearrange_str = f'{input_pattern} -> {output_pattern}'
+
+        rearrange_kwargs = {f'p{i}': p for i, p in enumerate(patch_size)}
+
+        self.to_patch_embedding = nn.Sequential(
+            Rearrange(rearrange_str, **rearrange_kwargs),
+            nn.Linear(patch_dim, dim),
+            nn.LayerNorm(dim),
+        )
+
+        self.dropout = nn.Dropout(emb_dropout)
+
+        # Create rotary embeddings
+        self.rotary_emb = GoldenGateRoPENd(
+            dim_pos = ndim,
+            heads = heads,
+            dim_head = dim_head,
+            rope_min_freq = rope_min_freq,
+            rope_max_freq = rope_max_freq,
+            rope_p_zero_freqs = rope_p_zero_freqs
+        )
+
+        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout, rotary_emb = self.rotary_emb)
+
+        self.to_latent = nn.Identity()
+        self.mlp_head = nn.Linear(dim, num_classes)
+
+    def muon_parameters(self):
+        params = []
+
+        for m in self.modules():
+            if isinstance(m, Attention):
+                params.extend([
+                    m.to_v.weight,
+                    m.to_out[0].weight
+                ])
+            elif isinstance(m, FeedForward):
+                params.extend([
+                    m.net[1].weight,
+                    m.net[-2].weight
+                ])
+
+        return params
+
+    def forward(
+        self,
+        x,
+        return_embed = False
+    ):
+        x = self.to_patch_embedding(x) # (b, *spatial_dims, patch_dim)
+
+        batch, *spatial_dims, _, device = *x.shape, x.device
+
+        # Generate position coordinates
+
+        grids = [arange(d, device = device, dtype = torch.float32) for d in spatial_dims]
+        grid = torch.meshgrid(*grids, indexing = 'ij')
+        pos = stack(grid, dim = -1)  # (*spatial_dims, ndim)
+
+        # flatten spatial dimensions for attention with nd rotary
+
+        pos = repeat(pos, '... p -> b (...) p', b = batch)
+        x, packed_shape = pack([x], 'b * d')
+
+        x = self.dropout(x)
+
+        embed = self.transformer(x, pos)
+
+        # return the embed with reconstituted patch shape
+
+        if return_embed:
+            embed, = unpack(embed, packed_shape, 'b * d')
+            return embed
+
+        # pooling to logits
+
+        pooled = reduce(embed, 'b n d -> b d', 'mean')
+
+        pooled = self.to_latent(pooled)
+        return self.mlp_head(pooled)
+
+
+if __name__ == '__main__':
+
+    model = ViTND(
+        ndim = 5,
+        input_shape = (4, 8, 16, 32, 64),
+        patch_size = (2, 2, 4, 4, 8),
+        num_classes = 1000,
+        dim = 512,
+        depth = 6,
+        heads = 8,
+        mlp_dim = 2048,
+        channels = 3,
+        dropout = 0.1,
+        emb_dropout = 0.1
+    )
+
+    data = torch.randn(2, 3, 4, 8, 16, 32, 64)
+
+    logits = model(data)
+
+    embed = model(data, return_embed = True) # (2, 2, 4, 4, 8, 8, 512)

vit_pytorch/vit_with_decorr.py

@@ -0,0 +1,298 @@
+# https://arxiv.org/abs/2510.14657
+# but instead of their decorr module updated with SGD, remove all projections and just return a decorrelation auxiliary loss
+
+import torch
+from torch import nn, stack, tensor
+import torch.nn.functional as F
+from torch.nn import Module, ModuleList
+
+from einops import rearrange, repeat, reduce, einsum, pack, unpack
+from einops.layers.torch import Rearrange
+
+# helpers
+
+def exists(v):
+    return v is not None
+
+def default(v, d):
+    return v if exists(v) else d
+
+def pair(t):
+    return t if isinstance(t, tuple) else (t, t)
+
+def divisible_by(num, den):
+    return (num % den) == 0
+
+# decorr loss
+
+class DecorrelationLoss(Module):
+    def __init__(
+        self,
+        sample_frac = 1.,
+        soft_validate_num_sampled = False,
+        use_subspace = False,
+        dim = None,
+        dim_subspace = 64,
+        num_subspaces = 1,
+        mean_center = False, # covariance if true
+        across_depth = False
+    ):
+        super().__init__()
+        assert 0. <= sample_frac <= 1.
+        self.need_sample = sample_frac < 1. and not across_depth
+        self.sample_frac = sample_frac
+
+        self.soft_validate_num_sampled = soft_validate_num_sampled
+        self.use_subspace = use_subspace
+        self.dim_subspace = dim_subspace
+        self.num_subspaces = num_subspaces
+        self.mean_center = mean_center
+
+        self.register_buffer('zero', tensor(0.), persistent = False)
+
+        if use_subspace:
+            assert exists(dim), 'dim must be passed in if using subspaces'
+            assert exists(dim_subspace) or num_subspaces == 1, 'dim_subspace must be provided if num_subspace greater than one'
+            assert dim_subspace < dim, 'subspace dimension must be less than or equal to feature dimension'
+
+            self.register_buffer('proj', torch.empty(num_subspaces, dim, dim_subspace), persistent = True)
+
+            for i in range(num_subspaces):
+                nn.init.orthogonal_(self.proj[i])
+
+    def forward(
+        self,
+        tokens
+    ):
+        batch, seq_len, dim, device = *tokens.shape[-3:], tokens.device
+
+        if self.need_sample:
+            num_sampled = int(seq_len * self.sample_frac)
+
+            assert self.soft_validate_num_sampled or num_sampled >= 2.
+
+            if num_sampled <= 1:
+                return self.zero
+
+            tokens, packed_shape = pack([tokens], '* n d e')
+
+            indices = torch.randn(tokens.shape[:2]).argsort(dim = -1)[..., :num_sampled, :]
+
+            batch_arange = torch.arange(tokens.shape[0], device = tokens.device)
+            batch_arange = rearrange(batch_arange, 'b -> b 1')
+
+            tokens = tokens[batch_arange, indices]
+            tokens, = unpack(tokens, packed_shape, '* n d e')
+
+        if self.use_subspace:
+            tokens = einsum(tokens, self.proj, '... n d, s d e -> ... s n e')
+            dim = self.dim_subspace
+        else:
+            tokens = rearrange(tokens, '... n d -> ... 1 n d')
+
+        if self.mean_center:
+            tokens = tokens - tokens.mean(dim = -2, keepdim = True)
+
+        dist = einsum(tokens, tokens, '... s n d, ... s n e -> ... s d e') / tokens.shape[-2]
+        eye = torch.eye(dim, device = device)
+
+        loss = dist.pow(2) * (1. - eye) / ((dim - 1) * dim)
+
+        loss = reduce(loss, '... b s d e -> b', 'sum')
+        return loss.mean()
+
+# classes
+
+class FeedForward(Module):
+    def __init__(self, dim, hidden_dim, dropout = 0.):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+
+        self.net = nn.Sequential(
+            nn.Linear(dim, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, dim),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x):
+        normed = self.norm(x)
+        return self.net(normed), normed
+
+class Attention(Module):
+    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):
+        super().__init__()
+        inner_dim = dim_head *  heads
+        project_out = not (heads == 1 and dim_head == dim)
+
+        self.norm = nn.LayerNorm(dim)
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+
+        self.attend = nn.Softmax(dim = -1)
+        self.dropout = nn.Dropout(dropout)
+
+        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, dim),
+            nn.Dropout(dropout)
+        ) if project_out else nn.Identity()
+
+    def forward(self, x):
+        normed = self.norm(x)
+
+        qkv = self.to_qkv(normed).chunk(3, dim = -1)
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
+
+        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+
+        attn = self.attend(dots)
+        attn = self.dropout(attn)
+
+        out = torch.matmul(attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+
+        return self.to_out(out), normed
+
+class Transformer(Module):
+    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0.):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.layers = ModuleList([])
+
+        for _ in range(depth):
+            self.layers.append(ModuleList([
+                Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout),
+                FeedForward(dim, mlp_dim, dropout = dropout)
+            ]))
+
+    def forward(self, x):
+
+        normed_inputs = []
+        layer_outputs = []
+
+        for attn, ff in self.layers:
+            attn_out, attn_normed_inp = attn(x)
+            layer_outputs.append(attn_out)
+            x = attn_out + x
+
+            ff_out, ff_normed_inp = ff(x)
+            layer_outputs.append(ff_out)
+            x = ff_out + x
+
+            normed_inputs.append(attn_normed_inp)
+            normed_inputs.append(ff_normed_inp)
+
+        return self.norm(x), stack(normed_inputs), stack(layer_outputs)
+
+class ViT(Module):
+    def __init__(
+        self,
+        *,
+        image_size,
+        patch_size,
+        num_classes,
+        dim,
+        depth,
+        heads,
+        mlp_dim,
+        pool = 'cls',
+        channels = 3,
+        dim_head = 64,
+        dropout = 0.,
+        emb_dropout = 0.,
+        decorr_sample_frac = 1.,
+        decorr_use_subspace = False,
+        decorr_dim_subspace = 64,
+        decorr_num_subspaces = 1,
+        decorr_mean_center = False,
+        decorr_layer_outputs_across_depth = False
+    ):
+        super().__init__()
+        image_height, image_width = pair(image_size)
+        self.patch_size = patch_height, patch_width = pair(patch_size)
+
+        assert divisible_by(image_height, patch_height) and divisible_by(image_width, patch_width), 'Image dimensions must be divisible by the patch size.'
+
+        num_patches = (image_height // patch_height) * (image_width // patch_width)
+        patch_dim = channels * patch_height * patch_width
+        assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'
+
+        self.to_patch_embedding = nn.Sequential(
+            Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_height, p2 = patch_width),
+            nn.LayerNorm(patch_dim),
+            nn.Linear(patch_dim, dim),
+            nn.LayerNorm(dim),
+        )
+
+        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
+        self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
+        self.dropout = nn.Dropout(emb_dropout)
+
+        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)
+
+        self.pool = pool
+        self.to_latent = nn.Identity()
+
+        self.mlp_head = nn.Linear(dim, num_classes)
+
+        # decorrelation loss related
+
+        self.decorr_layer_outputs_across_depth = decorr_layer_outputs_across_depth
+
+        self.has_decorr_loss = decorr_sample_frac > 0.
+
+        if self.has_decorr_loss:
+            self.decorr_loss = DecorrelationLoss(decorr_sample_frac, use_subspace = decorr_use_subspace, dim = dim, dim_subspace = decorr_dim_subspace, num_subspaces = decorr_num_subspaces, mean_center = decorr_mean_center, across_depth = self.decorr_layer_outputs_across_depth)
+
+        self.register_buffer('zero', torch.tensor(0.), persistent = False)
+
+    def forward(
+        self,
+        img,
+        return_decorr_aux_loss = None
+    ):
+        return_decorr_aux_loss = default(return_decorr_aux_loss, self.training) and self.has_decorr_loss
+
+        x = self.to_patch_embedding(img)
+        b, n, _ = x.shape
+
+        cls_tokens = repeat(self.cls_token, '1 1 d -> b 1 d', b = b)
+        x = torch.cat((cls_tokens, x), dim=1)
+        x += self.pos_embedding[:, :(n + 1)]
+        x = self.dropout(x)
+
+        x, normed_layer_inputs, layer_outputs = self.transformer(x)
+
+        # maybe return decor loss
+
+        decorr_aux_loss = self.zero
+
+        if return_decorr_aux_loss:
+            decorr_inputs = normed_layer_inputs
+
+            if self.decorr_layer_outputs_across_depth:
+                decorr_inputs = rearrange(layer_outputs, 'l b n d -> n b l d')
+
+            decorr_aux_loss = self.decorr_loss(decorr_inputs)
+
+        x = x.mean(dim = 1) if self.pool == 'mean' else x[:, 0]
+
+        x = self.to_latent(x)
+        return self.mlp_head(x), decorr_aux_loss
+
+# quick test
+
+if __name__ == '__main__':
+    decorr_loss = DecorrelationLoss(0.1)
+
+    hiddens = torch.randn(6, 2, 512, 256)
+
+    decorr_loss(hiddens)
+    decorr_loss(hiddens[0])
+
+    decorr_loss = DecorrelationLoss(0.0001, soft_validate_num_sampled = True)
+    out = decorr_loss(hiddens)
+    assert out.item() == 0

vit_pytorch/vit_with_keel_post_ln.py

@@ -0,0 +1,217 @@
+from __future__ import annotations
+
+import torch
+from torch import nn
+from torch.nn import Module, ModuleList
+
+from einops import rearrange, repeat
+from einops.layers.torch import Rearrange
+
+# functions
+
+def exists(v):
+    return v is not None
+
+def default(v, d):
+    return v if exists(v) else d
+
+def pair(t):
+    return t if isinstance(t, tuple) else (t, t)
+
+# classes
+
+class FeedForward(Module):
+    def __init__(self, dim, hidden_dim, dropout = 0.):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.LayerNorm(dim, bias = False),
+            nn.Linear(dim, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, dim),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+class Attention(Module):
+    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):
+        super().__init__()
+        inner_dim = dim_head *  heads
+        project_out = not (heads == 1 and dim_head == dim)
+
+        self.norm = nn.LayerNorm(dim, bias = False)
+
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+
+        self.attend = nn.Softmax(dim = -1)
+        self.dropout = nn.Dropout(dropout)
+
+        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, dim),
+            nn.Dropout(dropout)
+        ) if project_out else nn.Identity()
+
+    def forward(self, x):
+        x = self.norm(x)
+
+        qkv = self.to_qkv(x).chunk(3, dim = -1)
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
+
+        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+
+        attn = self.attend(dots)
+        attn = self.dropout(attn)
+
+        out = torch.matmul(attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)
+
+class Transformer(Module):
+    def __init__(
+        self,
+        dim,
+        depth,
+        heads,
+        dim_head,
+        mlp_dim,
+        dropout = 0.,
+        keel_residual_scale = None
+    ):
+        super().__init__()
+        assert depth > 1
+
+        self.layers = ModuleList([])
+
+        for _ in range(depth):
+            self.layers.extend([
+                Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout),
+                FeedForward(dim, mlp_dim, dropout = dropout)
+            ])
+
+        num_layers = depth * 2
+        self.keel_residual_scale = default(keel_residual_scale, num_layers)
+
+        self.post_norms = ModuleList([nn.LayerNorm(dim, bias = False) for _ in range(num_layers - 1)])
+
+    def forward(self, x):
+        residual_scale = self.keel_residual_scale
+
+        for layer_ind, layer in enumerate(self.layers):
+            first_layer = layer_ind == 0
+
+            residual = x
+
+            out = layer(x)
+
+            if first_layer:
+                x = out + residual
+                continue
+
+            post_norm = self.post_norms[layer_ind - 1]
+
+            x = post_norm(out + residual * residual_scale)
+
+        return x
+
+class ViT(Module):
+    def __init__(
+        self,
+        *,
+        image_size,
+        patch_size,
+        num_classes,
+        dim,
+        depth,
+        heads,
+        mlp_dim,
+        pool = 'cls',
+        channels = 3,
+        dim_head = 64,
+        dropout = 0.,
+        emb_dropout = 0.,
+        keel_residual_scale = None
+    ):
+        super().__init__()
+        image_height, image_width = pair(image_size)
+        self.patch_size = patch_height, patch_width = pair(patch_size)
+
+        assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
+
+        num_patches = (image_height // patch_height) * (image_width // patch_width)
+        patch_dim = channels * patch_height * patch_width
+
+        assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'
+        num_cls_tokens = 1 if pool == 'cls' else 0
+
+        self.to_patch_embedding = nn.Sequential(
+            Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_height, p2 = patch_width),
+            nn.LayerNorm(patch_dim),
+            nn.Linear(patch_dim, dim),
+            nn.LayerNorm(dim),
+        )
+
+        self.cls_token = nn.Parameter(torch.randn(num_cls_tokens, dim))
+        self.pos_embedding = nn.Parameter(torch.randn(num_patches + num_cls_tokens, dim))
+
+        self.dropout = nn.Dropout(emb_dropout)
+
+        self.transformer = Transformer(
+            dim,
+            depth,
+            heads,
+            dim_head,
+            mlp_dim,
+            dropout,
+            keel_residual_scale = keel_residual_scale
+        )
+
+        self.pool = pool
+        self.to_latent = nn.Identity()
+
+        self.mlp_head = nn.Linear(dim, num_classes) if num_classes > 0 else None
+
+    def forward(self, img):
+        batch = img.shape[0]
+        x = self.to_patch_embedding(img)
+
+        cls_tokens = repeat(self.cls_token, '... d -> b ... d', b = batch)
+        x = torch.cat((cls_tokens, x), dim = 1)
+
+        seq = x.shape[1]
+
+        x = x + self.pos_embedding[:seq]
+        x = self.dropout(x)
+
+        x = self.transformer(x)
+
+        if not exists(self.mlp_head):
+            return x
+
+        x = x.mean(dim = 1) if self.pool == 'mean' else x[:, 0]
+
+        x = self.to_latent(x)
+        return self.mlp_head(x)
+
+if __name__ == '__main__':
+    v = ViT(
+        image_size = 256,
+        patch_size = 32,
+        num_classes = 1000,
+        dim = 1024,
+        depth = 6,
+        heads = 16,
+        mlp_dim = 2048,
+        dropout = 0.1,
+        emb_dropout = 0.1
+    )
+
+    img = torch.randn(1, 3, 256, 256)
+
+    preds = v(img)
+
+    assert preds.shape == (1, 1000)

vit_pytorch/vit_with_patch_dropout.py

@@ -30,18 +30,11 @@ class PatchDropout(nn.Module):
 
         return x[batch_indices, patch_indices_keep]
 
-class PreNorm(nn.Module):
-    def __init__(self, dim, fn):
-        super().__init__()
-        self.norm = nn.LayerNorm(dim)
-        self.fn = fn
-    def forward(self, x, **kwargs):
-        return self.fn(self.norm(x), **kwargs)
-
 class FeedForward(nn.Module):
     def __init__(self, dim, hidden_dim, dropout = 0.):
         super().__init__()
         self.net = nn.Sequential(
+            nn.LayerNorm(dim),
             nn.Linear(dim, hidden_dim),
             nn.GELU(),
             nn.Dropout(dropout),
@@ -60,6 +53,7 @@ class Attention(nn.Module):
         self.heads = heads
         self.scale = dim_head ** -0.5
 
+        self.norm = nn.LayerNorm(dim)
         self.attend = nn.Softmax(dim = -1)
         self.dropout = nn.Dropout(dropout)
 
@@ -71,6 +65,7 @@ class Attention(nn.Module):
         ) if project_out else nn.Identity()
 
     def forward(self, x):
+        x = self.norm(x)
         qkv = self.to_qkv(x).chunk(3, dim = -1)
         q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
 
@@ -89,8 +84,8 @@ class Transformer(nn.Module):
         self.layers = nn.ModuleList([])
         for _ in range(depth):
             self.layers.append(nn.ModuleList([
-                PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout)),
-                PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout))
+                Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout),
+                FeedForward(dim, mlp_dim, dropout = dropout)
             ]))
     def forward(self, x):
         for attn, ff in self.layers:
@@ -102,7 +97,7 @@ class ViT(nn.Module):
     def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, pool = 'cls', channels = 3, dim_head = 64, dropout = 0., emb_dropout = 0., patch_dropout = 0.25):
         super().__init__()
         image_height, image_width = pair(image_size)
-        patch_height, patch_width = pair(patch_size)
+        self.patch_size = patch_height, patch_width = pair(patch_size)
 
         assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
 

vit_pytorch/vit_with_patch_merger.py

@@ -32,18 +32,11 @@ class PatchMerger(nn.Module):
 
 # classes
 
-class PreNorm(nn.Module):
-    def __init__(self, dim, fn):
-        super().__init__()
-        self.norm = nn.LayerNorm(dim)
-        self.fn = fn
-    def forward(self, x, **kwargs):
-        return self.fn(self.norm(x), **kwargs)
-
 class FeedForward(nn.Module):
     def __init__(self, dim, hidden_dim, dropout = 0.):
         super().__init__()
         self.net = nn.Sequential(
+            nn.LayerNorm(dim),
             nn.Linear(dim, hidden_dim),
             nn.GELU(),
             nn.Dropout(dropout),
@@ -62,6 +55,7 @@ class Attention(nn.Module):
         self.heads = heads
         self.scale = dim_head ** -0.5
 
+        self.norm = nn.LayerNorm(dim)
         self.attend = nn.Softmax(dim = -1)
         self.dropout = nn.Dropout(dropout)
 
@@ -73,6 +67,7 @@ class Attention(nn.Module):
         ) if project_out else nn.Identity()
 
     def forward(self, x):
+        x = self.norm(x)
         qkv = self.to_qkv(x).chunk(3, dim = -1)
         q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
 
@@ -88,6 +83,7 @@ class Attention(nn.Module):
 class Transformer(nn.Module):
     def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0., patch_merge_layer = None, patch_merge_num_tokens = 8):
         super().__init__()
+        self.norm = nn.LayerNorm(dim)
         self.layers = nn.ModuleList([])
 
         self.patch_merge_layer_index = default(patch_merge_layer, depth // 2) - 1 # default to mid-way through transformer, as shown in paper
@@ -95,8 +91,8 @@ class Transformer(nn.Module):
 
         for _ in range(depth):
             self.layers.append(nn.ModuleList([
-                PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout)),
-                PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout))
+                Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout),
+                FeedForward(dim, mlp_dim, dropout = dropout)
             ]))
     def forward(self, x):
         for index, (attn, ff) in enumerate(self.layers):
@@ -106,13 +102,13 @@ class Transformer(nn.Module):
             if index == self.patch_merge_layer_index:
                 x = self.patch_merger(x)
 
-        return x
+        return self.norm(x)
 
 class ViT(nn.Module):
     def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, patch_merge_layer = None, patch_merge_num_tokens = 8, channels = 3, dim_head = 64, dropout = 0., emb_dropout = 0.):
         super().__init__()
         image_height, image_width = pair(image_size)
-        patch_height, patch_width = pair(patch_size)
+        self.patch_size = patch_height, patch_width = pair(patch_size)
 
         assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
 
@@ -133,7 +129,6 @@ class ViT(nn.Module):
 
         self.mlp_head = nn.Sequential(
             Reduce('b n d -> b d', 'mean'),
-            nn.LayerNorm(dim),
             nn.Linear(dim, num_classes)
         )
 

vit_pytorch/vivit.py

@@ -1,5 +1,10 @@
+from collections import namedtuple
+
 import torch
-from torch import nn
+from torch import nn, cat
+import torch.nn.functional as F
+from torch.nn import Module, ModuleList
+from torch.nn.attention import SDPBackend, sdpa_kernel
 
 from einops import rearrange, repeat, reduce
 from einops.layers.torch import Rearrange
@@ -9,23 +14,19 @@ from einops.layers.torch import Rearrange
 def exists(val):
     return val is not None
 
+def divisible_by(num, den):
+    return (num % den) == 0
+
 def pair(t):
     return t if isinstance(t, tuple) else (t, t)
 
 # classes
 
-class PreNorm(nn.Module):
-    def __init__(self, dim, fn):
-        super().__init__()
-        self.norm = nn.LayerNorm(dim)
-        self.fn = fn
-    def forward(self, x, **kwargs):
-        return self.fn(self.norm(x), **kwargs)
-
-class FeedForward(nn.Module):
+class FeedForward(Module):
     def __init__(self, dim, hidden_dim, dropout = 0.):
         super().__init__()
         self.net = nn.Sequential(
+            nn.LayerNorm(dim),
             nn.Linear(dim, hidden_dim),
             nn.GELU(),
             nn.Dropout(dropout),
@@ -35,15 +36,18 @@ class FeedForward(nn.Module):
     def forward(self, x):
         return self.net(x)
 
-class Attention(nn.Module):
-    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):
+class Attention(Module):
+    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0., use_flash_attn = True):
         super().__init__()
+        self.use_flash_attn = use_flash_attn
+        self.dropout_p = dropout
         inner_dim = dim_head *  heads
         project_out = not (heads == 1 and dim_head == dim)
 
         self.heads = heads
         self.scale = dim_head ** -0.5
 
+        self.norm = nn.LayerNorm(dim)
         self.attend = nn.Softmax(dim = -1)
         self.dropout = nn.Dropout(dropout)
 
@@ -54,35 +58,100 @@ class Attention(nn.Module):
             nn.Dropout(dropout)
         ) if project_out else nn.Identity()
 
-    def forward(self, x):
+    def flash_attn(self, q, k, v, mask = None):
+
+        with sdpa_kernel([SDPBackend.MATH, SDPBackend.EFFICIENT_ATTENTION, SDPBackend.FLASH_ATTENTION, SDPBackend.CUDNN_ATTENTION]):
+
+            out = F.scaled_dot_product_attention(
+                q, k, v,
+                attn_mask = mask,
+                dropout_p = self.dropout_p,
+                is_causal = False,
+                scale = self.scale
+            )
+
+        return out
+
+    def forward(self, x, mask = None):
+        batch, seq, _ = x.shape
+
+        x = self.norm(x)
         qkv = self.to_qkv(x).chunk(3, dim = -1)
         q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
 
-        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+        if exists(mask):
+            mask = rearrange(mask, 'b j -> b 1 1 j')
 
-        attn = self.attend(dots)
-        attn = self.dropout(attn)
+        if self.use_flash_attn:
+            out =  self.flash_attn(q, k, v, mask = mask)
+
+        else:
+            dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+
+            if exists(mask):
+                dots = dots.masked_fill(~mask, -torch.finfo(dots.dtype).max)
+
+            attn = self.attend(dots)
+            attn = self.dropout(attn)
+
+            out = torch.matmul(attn, v)
 
-        out = torch.matmul(attn, v)
         out = rearrange(out, 'b h n d -> b n (h d)')
         return self.to_out(out)
 
-class Transformer(nn.Module):
-    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0.):
+class Transformer(Module):
+    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0., use_flash_attn = True):
         super().__init__()
-        self.layers = nn.ModuleList([])
+        self.use_flash_attn = use_flash_attn
+
+        self.norm = nn.LayerNorm(dim)
+        self.layers = ModuleList([])
         for _ in range(depth):
             self.layers.append(nn.ModuleList([
-                PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout)),
-                PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout))
+                Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout, use_flash_attn = use_flash_attn),
+                FeedForward(dim, mlp_dim, dropout = dropout)
             ]))
-    def forward(self, x):
-        for attn, ff in self.layers:
-            x = attn(x) + x
-            x = ff(x) + x
-        return x
 
-class ViT(nn.Module):
+    def forward(self, x, mask = None):
+
+        for attn, ff in self.layers:
+            x = attn(x, mask = mask) + x
+            x = ff(x) + x
+
+        return self.norm(x)
+
+class FactorizedTransformer(Module):
+    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0., use_flash_attn = True):
+        super().__init__()
+        self.use_flash_attn = use_flash_attn
+
+        self.norm = nn.LayerNorm(dim)
+        self.layers = nn.ModuleList([])
+
+        for _ in range(depth):
+            self.layers.append(nn.ModuleList([
+                Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout, use_flash_attn = use_flash_attn),
+                Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout, use_flash_attn = use_flash_attn),
+                FeedForward(dim, mlp_dim, dropout = dropout)
+            ]))
+
+    def forward(self, x, mask = None):
+        batch, frames, seq, _ = x.shape
+
+        if exists(mask):
+            mask = repeat(mask, 'b ... -> (b space) ...', space = x.shape[2])
+
+        for spatial_attn, temporal_attn, ff in self.layers:
+            x = rearrange(x, 'b f n d -> (b f) n d')
+            x = spatial_attn(x) + x
+            x = rearrange(x, '(b f) n d -> (b n) f d', b = batch, f = frames)
+            x = temporal_attn(x, mask = mask) + x
+            x = ff(x) + x
+            x = rearrange(x, '(b n) f d -> b f n d', b = batch, n = seq)
+
+        return self.norm(x)
+
+class ViViT(Module):
     def __init__(
         self,
         *,
@@ -100,14 +169,17 @@ class ViT(nn.Module):
         channels = 3,
         dim_head = 64,
         dropout = 0.,
-        emb_dropout = 0.
+        emb_dropout = 0.,
+        variant = 'factorized_encoder',
+        use_flash_attn: bool = True,
     ):
         super().__init__()
         image_height, image_width = pair(image_size)
-        patch_height, patch_width = pair(image_patch_size)
+        self.patch_size = patch_height, patch_width = pair(image_patch_size)
 
-        assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
-        assert frames % frame_patch_size == 0, 'Frames must be divisible by frame patch size'
+        assert divisible_by(image_height, patch_height) and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
+        assert divisible_by(frames, frame_patch_size), 'Frames must be divisible by frame patch size'
+        assert variant in ('factorized_encoder', 'factorized_self_attention'), f'variant = {variant} is not implemented'
 
         num_image_patches = (image_height // patch_height) * (image_width // patch_width)
         num_frame_patches = (frames // frame_patch_size)
@@ -116,10 +188,12 @@ class ViT(nn.Module):
 
         assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'
 
+        self.frame_patch_size = frame_patch_size
+
         self.global_average_pool = pool == 'mean'
 
         self.to_patch_embedding = nn.Sequential(
-            Rearrange('b c (f pf) (h p1) (w p2) -> b f (h w) (p1 p2 pf c)', p1 = patch_height, p2 = patch_width, pf = frame_patch_size),
+            Rearrange('b c (f pf) (h p1) (w p2) -> b f (h w) (pf p1 p2 c)', p1 = patch_height, p2 = patch_width, pf = frame_patch_size),
             nn.LayerNorm(patch_dim),
             nn.Linear(patch_dim, dim),
             nn.LayerNorm(dim)
@@ -129,57 +203,103 @@ class ViT(nn.Module):
         self.dropout = nn.Dropout(emb_dropout)
 
         self.spatial_cls_token = nn.Parameter(torch.randn(1, 1, dim)) if not self.global_average_pool else None
-        self.temporal_cls_token = nn.Parameter(torch.randn(1, 1, dim)) if not self.global_average_pool else None
 
-        self.spatial_transformer = Transformer(dim, spatial_depth, heads, dim_head, mlp_dim, dropout)
-        self.temporal_transformer = Transformer(dim, temporal_depth, heads, dim_head, mlp_dim, dropout)
+        if variant == 'factorized_encoder':
+            self.temporal_cls_token = nn.Parameter(torch.randn(1, 1, dim)) if not self.global_average_pool else None
+            self.spatial_transformer = Transformer(dim, spatial_depth, heads, dim_head, mlp_dim, dropout, use_flash_attn)
+            self.temporal_transformer = Transformer(dim, temporal_depth, heads, dim_head, mlp_dim, dropout, use_flash_attn)
+        elif variant == 'factorized_self_attention':
+            assert spatial_depth == temporal_depth, 'Spatial and temporal depth must be the same for factorized self-attention'
+            self.factorized_transformer = FactorizedTransformer(dim, spatial_depth, heads, dim_head, mlp_dim, dropout, use_flash_attn)
 
         self.pool = pool
         self.to_latent = nn.Identity()
 
-        self.mlp_head = nn.Sequential(
-            nn.LayerNorm(dim),
-            nn.Linear(dim, num_classes)
-        )
+        self.mlp_head = nn.Linear(dim, num_classes)
+        self.variant = variant
+
+    def forward(self, video, mask = None):
+        device = video.device
 
-    def forward(self, video):
         x = self.to_patch_embedding(video)
-        b, f, n, _ = x.shape
+        batch, frames, seq, _ = x.shape
 
-        x = x + self.pos_embedding[:, :f, :n]
+        x = x + self.pos_embedding[:, :frames, :seq]
 
         if exists(self.spatial_cls_token):
-            spatial_cls_tokens = repeat(self.spatial_cls_token, '1 1 d -> b f 1 d', b = b, f = f)
-            x = torch.cat((spatial_cls_tokens, x), dim = 2)
+            spatial_cls_tokens = repeat(self.spatial_cls_token, '1 1 d -> b f 1 d', b = batch, f = frames)
+            x = cat((spatial_cls_tokens, x), dim = 2)
 
         x = self.dropout(x)
 
-        x = rearrange(x, 'b f n d -> (b f) n d')
+        # maybe temporal mask
 
-        # attend across space
+        temporal_mask = None
 
-        x = self.spatial_transformer(x)
+        if exists(mask):
+            temporal_mask = reduce(mask, 'b (f patch) -> b f', 'all', patch = self.frame_patch_size)
 
-        x = rearrange(x, '(b f) n d -> b f n d', b = b)
+        # the two variants
 
-        # excise out the spatial cls tokens or average pool for temporal attention
+        if self.variant == 'factorized_encoder':
+            x = rearrange(x, 'b f n d -> (b f) n d')
 
-        x = x[:, :, 0] if not self.global_average_pool else reduce(x, 'b f n d -> b f d', 'mean')
+            # attend across space
 
-        # append temporal CLS tokens
+            x = self.spatial_transformer(x)
+            x = rearrange(x, '(b f) n d -> b f n d', b = batch)
 
-        if exists(self.temporal_cls_token):
-            temporal_cls_tokens = repeat(self.temporal_cls_token, '1 1 d-> b 1 d', b = b)
+            # excise out the spatial cls tokens or average pool for temporal attention
 
-            x = torch.cat((temporal_cls_tokens, x), dim = 1)
+            x = x[:, :, 0] if not self.global_average_pool else reduce(x, 'b f n d -> b f d', 'mean')
 
-        # attend across time
+            # append temporal CLS tokens
 
-        x = self.temporal_transformer(x)
+            if exists(self.temporal_cls_token):
+                temporal_cls_tokens = repeat(self.temporal_cls_token, '1 1 d-> b 1 d', b = batch)
 
-        # excise out temporal cls token or average pool
+                x = cat((temporal_cls_tokens, x), dim = 1)
 
-        x = x[:, 0] if not self.global_average_pool else reduce(x, 'b f d -> b d', 'mean')
+                if exists(temporal_mask):
+                    temporal_mask = F.pad(temporal_mask, (1, 0), value = True)
+
+            # attend across time
+
+            x = self.temporal_transformer(x, mask = temporal_mask)
+
+            # excise out temporal cls token or average pool
+
+            x = x[:, 0] if not self.global_average_pool else reduce(x, 'b f d -> b d', 'mean')
+
+        elif self.variant == 'factorized_self_attention':
+
+            x = self.factorized_transformer(x, mask = temporal_mask)
+
+            x = x[:, 0, 0] if not self.global_average_pool else reduce(x, 'b f n d -> b d', 'mean')
 
         x = self.to_latent(x)
         return self.mlp_head(x)
+
+# main
+
+if __name__ == '__main__':
+
+    vivit = ViViT(
+        dim = 512,
+        spatial_depth = 2,
+        temporal_depth = 2,
+        heads = 4,
+        mlp_dim = 2048,
+        image_size = 256,
+        image_patch_size = 16,
+        frames = 8,
+        frame_patch_size = 2,
+        num_classes = 1000,
+        variant = 'factorized_encoder',
+    )
+
+    video = torch.randn(3, 3, 8, 256, 256)
+    mask = torch.randint(0, 2, (3, 8)).bool()
+
+    logits = vivit(video, mask = None)
+    assert logits.shape == (3, 1000)

vit_pytorch/vivit_with_moss.py

@@ -0,0 +1,452 @@
+# https://openreview.net/forum?id=Co6SCyBIjo
+# applied at https://arxiv.org/abs/2605.03269 - 50-85% jump in pick-place moving conveyer belt
+
+import torch
+from torch import nn
+from torch.nn import Module, ModuleList
+import torch.nn.functional as F
+from torch.nn.attention import SDPBackend, sdpa_kernel
+
+from einops import rearrange, repeat, reduce, einsum
+from einops.layers.torch import Rearrange
+
+# helpers
+
+def exists(val):
+    return val is not None
+
+def default(val, d):
+    return val if exists(val) else d
+
+
+def divisible_by(num, den):
+    return (num % den) == 0
+
+def is_odd(n):
+    return not divisible_by(n, 2)
+
+def pair(t):
+    return t if isinstance(t, tuple) else (t, t)
+
+def l2norm(t):
+    return F.normalize(t, dim = -1)
+
+# normalization helpers
+
+class ChanLayerNorm(Module):
+    def __init__(self, dim, eps = 1e-5):
+        super().__init__()
+        self.eps = eps
+        self.gamma = nn.Parameter(torch.ones(1, dim, 1, 1))
+
+    def forward(self, x):
+        var = torch.var(x, dim = 1, unbiased = False, keepdim = True)
+        mean = torch.mean(x, dim = 1, keepdim = True)
+        return (x - mean) * (var + self.eps).rsqrt() * self.gamma
+
+# classes
+
+class FeedForward(Module):
+    def __init__(self, dim, hidden_dim, dropout = 0.):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Linear(dim, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, dim),
+            nn.Dropout(dropout)
+        )
+    def forward(self, x):
+        return self.net(x)
+
+class Attention(Module):
+    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0., use_flash_attn = True, causal = False):
+        super().__init__()
+        self.use_flash_attn = use_flash_attn
+        self.dropout_p = dropout
+        self.causal = causal
+        inner_dim = dim_head * heads
+        project_out = not (heads == 1 and dim_head == dim)
+
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+
+        self.norm = nn.LayerNorm(dim)
+        self.attend = nn.Softmax(dim = -1)
+        self.dropout = nn.Dropout(dropout)
+
+        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, dim),
+            nn.Dropout(dropout)
+        ) if project_out else nn.Identity()
+
+    def flash_attn(self, q, k, v, mask = None):
+        is_causal = self.causal and q.shape[-2] > 1
+
+        with sdpa_kernel([SDPBackend.MATH, SDPBackend.EFFICIENT_ATTENTION, SDPBackend.FLASH_ATTENTION, SDPBackend.CUDNN_ATTENTION]):
+            out = F.scaled_dot_product_attention(
+                q, k, v,
+                attn_mask = mask,
+                dropout_p = self.dropout_p,
+                is_causal = is_causal,
+                scale = self.scale
+            )
+        return out
+
+    def forward(self, x, mask = None, cache = None, return_cache = False):
+        is_causal = self.causal and x.shape[-2] > 1
+        assert not (is_causal and exists(mask)), 'causal attention is not compatible with key padding mask'
+
+        x = self.norm(x)
+        qkv = self.to_qkv(x).chunk(3, dim = -1)
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
+
+        if exists(cache):
+            ck, cv = cache
+            k = torch.cat((ck, k), dim = -2)
+            v = torch.cat((cv, v), dim = -2)
+
+        if exists(mask):
+            mask = rearrange(mask, 'b j -> b 1 1 j')
+
+        if self.use_flash_attn:
+            out = self.flash_attn(q, k, v, mask = mask)
+        else:
+            dots = einsum(q, k, 'b h i d, b h j d -> b h i j') * self.scale
+            if exists(mask):
+                dots = dots.masked_fill(~mask, -torch.finfo(dots.dtype).max)
+
+            if self.causal:
+                i, j = dots.shape[-2:]
+                causal_mask = torch.ones((i, j), device = x.device, dtype = torch.bool).triu(j - i + 1)
+                dots = dots.masked_fill(causal_mask, -torch.finfo(dots.dtype).max)
+
+            attn = self.attend(dots)
+            attn = self.dropout(attn)
+            out = torch.matmul(attn, v)
+
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        out = self.to_out(out)
+
+        if not return_cache:
+            return out
+
+        return out, (k, v)
+
+class Transformer(Module):
+    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0., use_flash_attn = True, causal = False):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.layers = ModuleList([])
+
+        for _ in range(depth):
+            self.layers.append(ModuleList([
+                Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout, use_flash_attn = use_flash_attn, causal = causal),
+                FeedForward(dim, mlp_dim, dropout = dropout)
+            ]))
+
+    def forward(self, x, mask = None, cache = None, return_cache = False):
+        new_caches = []
+        cache = default(cache, (None,) * len(self.layers))
+
+        for (attn, ff), layer_cache in zip(self.layers, cache):
+            attn_out, next_cache = attn(x, mask = mask, cache = layer_cache, return_cache = True)
+            new_caches.append(next_cache)
+
+            x = attn_out + x
+            x = ff(x) + x
+
+        x = self.norm(x)
+
+        if not return_cache:
+            return x
+
+        return x, tuple(new_caches)
+
+# moss specific classes
+
+class STSSEncoder(Module):
+    def __init__(self, dim, local_time = 3, local_height = 3, local_width = 3, hidden_dim = 64):
+        super().__init__()
+
+        self.spatial_to_hidden = nn.Linear(local_height * local_width, hidden_dim)
+
+        self.conv = nn.Sequential(
+            nn.Conv2d(hidden_dim, hidden_dim, 3, padding = 1),
+            ChanLayerNorm(hidden_dim),
+            nn.GELU(),
+            nn.Conv2d(hidden_dim, hidden_dim, 3, padding = 1),
+            ChanLayerNorm(hidden_dim),
+            nn.GELU()
+        )
+
+        self.time_to_out = nn.Linear(local_time * hidden_dim, dim)
+
+    def forward(self, sim):
+        b, t, h, w, lt, lh, lw = sim.shape
+
+        x = rearrange(sim, 'b t h w lt lh lw -> b t h w lt (lh lw)')
+        x = self.spatial_to_hidden(x)
+
+        x = rearrange(x, 'b t h w lt d -> (b t lt) d h w')
+        x = self.conv(x)
+        x = rearrange(x, '(b t lt) d h w -> b t h w (lt d)', b = b, t = t, lt = lt)
+
+        return self.time_to_out(x)
+
+class MOSS(Module):
+    def __init__(
+        self,
+        dim,
+        local_time = 3,
+        local_height = 3,
+        local_width = 3,
+        hidden_dim = 64,
+        orders = 2,
+        causal = False
+    ):
+        super().__init__()
+        assert is_odd(local_time) and is_odd(local_height) and is_odd(local_width), 'MOSS local dimensions must be odd'
+
+        self.local_time = local_time
+        self.local_height = local_height
+        self.local_width = local_width
+        self.causal = causal
+
+        self.encoders = ModuleList([STSSEncoder(dim, local_time, local_height, local_width, hidden_dim) for _ in range(orders)])
+        self.to_order_out = ModuleList([nn.Linear(dim, dim) for _ in range(orders)])
+        self.to_out = nn.Linear(dim, dim)
+
+    def stss_transform(self, x, cache = None, return_cache = False):
+        assert not (exists(cache) and not self.causal), 'cache cannot be passed in if MOSS is not causal'
+
+        lt, lh, lw = self.local_time, self.local_height, self.local_width
+        _, _, h, w, _ = x.shape
+
+        x = l2norm(x)
+        x = rearrange(x, 'b t h w c -> b c t h w')
+
+        pad_h, pad_w = lh // 2, lw // 2
+        pad_t_past, pad_t_future = (lt - 1, 0) if self.causal else (lt // 2, lt // 2)
+
+        has_cache = self.causal and exists(cache)
+        x_temporal = torch.cat((cache, x), dim = 2) if has_cache else x
+
+        padding = (pad_w, pad_w, pad_h, pad_h, 0 if has_cache else pad_t_past, pad_t_future)
+        padded_x = F.pad(x_temporal, padding)
+
+        windows = padded_x.unfold(2, lt, 1).unfold(3, lh, 1).unfold(4, lw, 1)
+
+        sim = einsum(x, windows, 'b c t h w, b c t h w l u v -> b t h w l u v')
+
+        if not return_cache:
+            return sim
+
+        new_cache = padded_x[..., -(lt - 1):, pad_h:(pad_h + h), pad_w:(pad_w + w)] if self.causal else None
+        return sim, new_cache
+
+    def forward(
+        self,
+        x,
+        cache = None,
+        return_cache = False
+    ):
+        assert not (exists(cache) and not self.causal), 'cache cannot be passed in if MOSS is not causal'
+
+        out = self.to_out(x)
+
+        new_caches = []
+        cache = default(cache, (None,) * len(self.encoders))
+
+        for encoder, to_order_out, layer_cache in zip(self.encoders, self.to_order_out, cache):
+            sim, next_cache = self.stss_transform(x, cache = layer_cache, return_cache = True)
+            new_caches.append(next_cache)
+
+            x = encoder(sim)
+            out = out + to_order_out(x)
+
+        if not return_cache:
+            return out
+
+        return out, tuple(new_caches)
+
+# main architecture
+
+class ViViT(Module):
+    def __init__(
+        self,
+        *,
+        image_size,
+        image_patch_size,
+        frames,
+        frame_patch_size,
+        num_classes,
+        dim,
+        spatial_depth,
+        temporal_depth,
+        heads,
+        mlp_dim,
+        pool = 'cls',
+        channels = 3,
+        dim_head = 64,
+        dropout = 0.,
+        emb_dropout = 0.,
+        use_flash_attn: bool = True,
+        moss_local_time = 3,
+        moss_local_height = 3,
+        moss_local_width = 3,
+        moss_hidden_dim = 64,
+        moss_orders = 2,
+        moss_causal = True,
+    ):
+        super().__init__()
+        image_height, image_width = pair(image_size)
+        self.patch_size = patch_height, patch_width = pair(image_patch_size)
+
+        assert divisible_by(image_height, patch_height) and divisible_by(image_width, patch_width), 'Image dimensions must be divisible by the patch size.'
+        assert divisible_by(frames, frame_patch_size), 'Frames must be divisible by frame patch size'
+        assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'
+
+        num_image_patches = (image_height // patch_height) * (image_width // patch_width)
+        num_frame_patches = frames // frame_patch_size
+        patch_dim = channels * patch_height * patch_width * frame_patch_size
+
+        self.frame_patch_size = frame_patch_size
+        self.patch_h = image_height // patch_height
+        self.patch_w = image_width // patch_width
+        self.global_average_pool = pool == 'mean'
+
+        self.to_patch_embedding = nn.Sequential(
+            Rearrange('b c (f pf) (h p1) (w p2) -> b f (h w) (pf p1 p2 c)', p1 = patch_height, p2 = patch_width, pf = frame_patch_size),
+            nn.LayerNorm(patch_dim),
+            nn.Linear(patch_dim, dim),
+            nn.LayerNorm(dim)
+        )
+
+        self.pos_embedding = nn.Parameter(torch.randn(1, num_frame_patches, num_image_patches, dim))
+        self.dropout = nn.Dropout(emb_dropout)
+
+        self.has_cls = not self.global_average_pool
+        self.spatial_cls_token = nn.Parameter(torch.randn(1, 1, dim)) if self.has_cls else None
+        self.temporal_cls_token = nn.Parameter(torch.randn(1, 1, dim)) if self.has_cls else None
+
+        self.spatial_transformer = Transformer(dim, spatial_depth, heads, dim_head, mlp_dim, dropout, use_flash_attn, causal = False)
+        self.temporal_transformer = Transformer(dim, temporal_depth, heads, dim_head, mlp_dim, dropout, use_flash_attn, causal = moss_causal)
+
+        self.moss = MOSS(
+            dim,
+            local_time = moss_local_time,
+            local_height = moss_local_height,
+            local_width = moss_local_width,
+            hidden_dim = moss_hidden_dim,
+            orders = moss_orders,
+            causal = moss_causal
+        )
+
+        self.mlp_head = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Linear(dim, num_classes)
+        )
+
+    def forward(self, video, mask = None):
+        assert not (exists(mask) and self.moss.causal), 'mask cannot be passed if MOSS is causal'
+
+        x = self.to_patch_embedding(video)
+        batch, frames, seq, _ = x.shape
+
+        x = x + self.pos_embedding[:, :frames, :seq]
+
+        if self.has_cls:
+            spatial_cls_tokens = repeat(self.spatial_cls_token, '1 1 d -> b f 1 d', b = batch, f = frames)
+            x = torch.cat((spatial_cls_tokens, x), dim = 2)
+
+        x = self.dropout(x)
+
+        # temporal mask
+
+        temporal_mask = None
+        if exists(mask):
+            temporal_mask = reduce(mask, 'b (f patch) -> b f', 'all', patch = self.frame_patch_size)
+
+        x = rearrange(x, 'b f n d -> (b f) n d')
+
+        # attend across space
+
+        x = self.spatial_transformer(x)
+        x = rearrange(x, '(b f) n d -> b f n d', b = batch)
+
+        # moss integration over spatial patch tokens
+
+        if self.has_cls:
+            spatial_cls_tokens, patch_tokens = x[:, :, :1], x[:, :, 1:]
+        else:
+            patch_tokens = x
+
+        patch_tokens = rearrange(patch_tokens, 'b f (h w) d -> b f h w d', h = self.patch_h, w = self.patch_w)
+        patch_tokens = self.moss(patch_tokens)
+        patch_tokens = rearrange(patch_tokens, 'b f h w d -> b f (h w) d')
+
+        # pool spatial features
+
+        moss_pooled = reduce(patch_tokens, 'b f n d -> b f d', 'mean')
+
+        if self.has_cls:
+            x = rearrange(spatial_cls_tokens, 'b f 1 d -> b f d') + moss_pooled
+        else:
+            x = moss_pooled
+
+        # append temporal cls tokens
+
+        if self.has_cls:
+            temporal_cls_tokens = repeat(self.temporal_cls_token, '1 1 d -> b 1 d', b = batch)
+            x = torch.cat((temporal_cls_tokens, x), dim = 1)
+
+            if exists(temporal_mask):
+                temporal_mask = F.pad(temporal_mask, (1, 0), value = True)
+
+        # attend across time
+
+        x = self.temporal_transformer(x, mask = temporal_mask)
+
+        # temporal pooling
+
+        x = x[:, 0] if self.has_cls else reduce(x, 'b f d -> b d', 'mean')
+
+        return self.mlp_head(x)
+
+if __name__ == '__main__':
+
+    vivit = ViViT(
+        dim = 512,
+        spatial_depth = 2,
+        temporal_depth = 2,
+        heads = 4,
+        mlp_dim = 2048,
+        image_size = 256,
+        image_patch_size = 32,
+        frames = 8,
+        frame_patch_size = 2,
+        num_classes = 1000,
+        moss_causal = True
+    )
+
+    video = torch.randn(2, 3, 8, 256, 256)
+    logits = vivit(video, mask = None)
+    assert logits.shape == (2, 1000)
+
+    moss = MOSS(
+        dim = 512,
+        local_time = 3,
+        local_height = 3,
+        local_width = 3,
+        hidden_dim = 64,
+        orders = 2,
+        causal = True
+    )
+
+    moss_input = torch.randn(2, 8, 16, 16, 512) # (batch, frames, height, width, dim)
+    moss_output = moss(moss_input)
+    assert moss_output.shape == (2, 8, 16, 16, 512)

vit_pytorch/wwt.py

@@ -0,0 +1,556 @@
+from __future__ import annotations
+
+from functools import partial
+from collections import namedtuple
+
+import torch
+from torch import nn, cat, is_tensor
+from torch.nn import Module, ModuleList, ParameterList
+from einops import rearrange, repeat, reduce, einsum, pack, unpack
+from einops.layers.torch import Rearrange
+
+from vit_pytorch.lejepa import sigreg_loss
+
+# constants
+
+LinearNoBias = partial(nn.Linear, bias = False)
+LayerNormNoBias = partial(nn.LayerNorm, bias = False)
+
+WWTReturn = namedtuple('WWTReturn', ['slot_logits', 'token_logits'])
+WWTFeatureReturn = namedtuple('WWTFeatureReturn', ['slots', 'tokens', 'masks'])
+
+# functions
+
+def exists(val):
+    return val is not None
+
+def default(val, d):
+    return val if exists(val) else d
+
+def pair(t):
+    return t if isinstance(t, tuple) else (t, t)
+
+def cast_tuple(val, length = 1):
+    return val if isinstance(val, tuple) else ((val,) * length)
+
+def is_empty(t):
+    return len(t) == 0
+
+def divisible_by(num, den):
+    return (num % den) == 0
+
+def l1norm(t, dim = -1, eps = 1e-8):
+    return t / t.sum(dim = dim, keepdim = True).clamp(min = eps)
+
+def pack_with_inverse(t, pattern):
+    packed, packed_shape = pack(t, pattern)
+
+    def inverse(packed_out):
+        return unpack(packed_out, packed_shape, pattern)
+
+    return packed, inverse
+
+# classes
+
+class AutoencodingHead(Module):
+    def __init__(
+        self,
+        *,
+        image_size,
+        patch_size,
+        decoder: Module | None = None,
+        pathways: tuple[tuple[int, ...], ...] | None = None,
+        patch_pathway_id = 0,
+        channel_first = False
+    ):
+        super().__init__()
+        image_height, image_width = pair(image_size)
+        patch_height, patch_width = pair(patch_size)
+
+        grid_h = image_height // patch_height
+        grid_w = image_width // patch_width
+
+        self.patch_pathway_id = patch_pathway_id
+
+        # customary pathways (e.g. 3 -> 2 -> 0 or 0 -> 2 -> 3)
+        # allows for enforcing consistency loss between different hierarchical traversals
+
+        if exists(pathways):
+            for pathway in pathways:
+                is_descending = pathway[0] > pathway[-1]
+                for source, target in zip(pathway[:-1], pathway[1:]):
+                    assert (target < source) == is_descending, 'pathway must be strictly directional'
+
+        self.pathways = pathways
+
+        pattern = 'b (h w) d -> b d h w' if channel_first else 'b (h w) d -> b h w d'
+        self.to_spatial = Rearrange(pattern, h = grid_h, w = grid_w)
+
+        self.decoder = default(decoder, nn.Identity())
+
+    def forward(self, hierarchy_features, masks, interactions):
+
+        # reduce masks
+
+        masks = {
+            interaction: reduce(m, 'b ... t s -> b t s', 'mean')
+            for m, interaction in zip(masks, interactions)
+        }
+
+        # default pathways
+
+        pathways = default(self.pathways, tuple((j, self.patch_pathway_id) for i, j in interactions if i == self.patch_pathway_id))
+        assert not is_empty(pathways), 'no valid pathways found'
+
+        # feature map construction
+
+        def construct_feature_map(pathway):
+            start, end = pathway[0], pathway[-1]
+            is_descending = start > end
+            features = hierarchy_features[start]
+
+            # chain interactions
+
+            for source, target in zip(pathway[:-1], pathway[1:]):
+                interaction = (target, source) if is_descending else (source, target)
+                assert interaction in masks, f'interaction {interaction} is missing'
+
+                mask = masks[interaction]
+
+                if not is_descending:
+                    mask = rearrange(mask, 'b i j -> b j i')
+
+                attn = mask.softmax(dim = -1)
+                features = einsum(attn, features, 'b t s, b s d -> b t d')
+
+            # rearrange spatial if ends in patches
+
+            if end == self.patch_pathway_id:
+                features = self.to_spatial(features)
+
+            return self.decoder(features)
+
+        feature_maps = tuple(construct_feature_map(p) for p in pathways)
+
+        return feature_maps[0] if len(feature_maps) == 1 else feature_maps
+
+def FeedForward(dim, hidden_dim, dropout = 0., out_dim = None):
+    return nn.Sequential(
+        LayerNormNoBias(dim),
+        nn.Linear(dim, hidden_dim),
+        nn.GELU(),
+        nn.Dropout(dropout),
+        nn.Linear(hidden_dim, default(out_dim, dim)),
+        nn.Dropout(dropout)
+    )
+
+class MutualAttention(Module):
+    def __init__(
+        self,
+        dim,
+        num_slots,
+        heads,
+        dim_head,
+        mlp_dim,
+        dropout = 0.,
+        l1norm_after_tokens_softmax = False,
+        token_softmax_over_slots = False,
+        project_mask_groups = False,
+        gate_attention = True
+    ):
+        super().__init__()
+        self.heads = heads
+        self.l1norm_after_tokens_softmax = l1norm_after_tokens_softmax
+        self.token_softmax_over_slots = token_softmax_over_slots
+
+        self.q_groups = 2 if token_softmax_over_slots else 1
+
+        self.project_mask_groups = project_mask_groups and token_softmax_over_slots
+        self.mask_groups = 1 if self.project_mask_groups else self.q_groups
+
+        inner_dim = heads * dim_head
+
+        self.to_q_v_tokens = LinearNoBias(dim, inner_dim * (self.q_groups + 1))
+        self.to_k_v_slots = LinearNoBias(dim, inner_dim * 2)
+
+        self.scale = dim_head ** -0.5
+
+        self.to_out_tokens = nn.Sequential(nn.Linear(inner_dim, dim), nn.Dropout(dropout))
+        self.to_out_slots = nn.Sequential(nn.Linear(inner_dim, dim), nn.Dropout(dropout))
+
+        if self.project_mask_groups:
+            self.mask_project = nn.Conv2d(self.q_groups * heads, heads, 1)
+
+        self.mlp_mask = FeedForward(
+            self.mask_groups * heads * num_slots + dim,
+            mlp_dim,
+            dropout = dropout,
+            out_dim = self.mask_groups * heads * num_slots
+        )
+
+        self.gate_attention = gate_attention
+        if gate_attention:
+            self.tokens_gate = nn.Sequential(
+                LinearNoBias(dim, heads),
+                Rearrange('b t h -> b h t 1'),
+                nn.Sigmoid()
+            )
+            self.slots_gate = nn.Sequential(
+                LinearNoBias(dim, heads),
+                Rearrange('b s h -> b h s 1'),
+                nn.Sigmoid()
+            )
+
+    def forward(self, tokens, slots, mask):
+        h, g = self.heads, self.q_groups
+
+        # queries, keys, values
+
+        tokens_q_v = rearrange(self.to_q_v_tokens(tokens), 'b t (g h d) -> g b h t d', h = h, g = g + 1)
+        q, v_tokens = tokens_q_v[:-1], tokens_q_v[-1]
+        q = rearrange(q, 'g b h t d -> b g h t d')
+
+        k, v_slots = rearrange(self.to_k_v_slots(slots), 'b s (kv h d) -> kv b h s d', h = h, kv = 2)
+
+        # mutual attention
+
+        sim = einsum(q, k, 'b g h t d, b h s d -> b g h t s') * self.scale
+        mask_prime = mask + sim
+
+        if self.token_softmax_over_slots:
+            mask_prime_slots, mask_prime_tokens = mask_prime.unbind(dim = 1)
+            attn_tokens = mask_prime_tokens.softmax(dim = -1)
+        else:
+            mask_prime_slots = rearrange(mask_prime, 'b 1 h t s -> b h t s')
+            attn_tokens = mask_prime_slots.softmax(dim = -2)
+
+        # slot softmax - https://arxiv.org/abs/2006.15055
+
+        attn_slots = mask_prime_slots.softmax(dim = -1)
+
+        if self.l1norm_after_tokens_softmax:
+            attn_slots = l1norm(attn_slots, dim = -2)
+
+        # aggregate
+
+        tokens_out_pre = einsum(attn_tokens, v_slots, 'b h t s, b h s d -> b h t d')
+        slots_out_pre = einsum(attn_slots, v_tokens, 'b h t s, b h t d -> b h s d')
+
+        if self.gate_attention:
+            tokens_out_pre = tokens_out_pre * self.tokens_gate(tokens)
+            slots_out_pre = slots_out_pre * self.slots_gate(slots)
+
+        tokens_out = self.to_out_tokens(rearrange(tokens_out_pre, 'b h t d -> b t (h d)'))
+        slots_out = self.to_out_slots(rearrange(slots_out_pre, 'b h s d -> b s (h d)'))
+
+        # mask update
+
+        if self.project_mask_groups:
+            mask_prime = self.mask_project(rearrange(mask_prime, 'b g h t s -> b (g h) t s'))
+            mask_prime = rearrange(mask_prime, 'b h t s -> b 1 h t s')
+
+        mask_next_reshaped = self.mlp_mask(cat((rearrange(mask_prime, 'b g h t s -> b t (g h s)'), tokens + tokens_out), dim = -1))
+        mask_next = rearrange(mask_next_reshaped, 'b t (g h s) -> b g h t s', h = h, g = self.mask_groups)
+
+        return tokens_out, slots_out, mask_next
+
+class WWTBlock(Module):
+    def __init__(
+        self,
+        dim,
+        num_hierarchies,
+        seq_lengths,
+        interactions,
+        heads,
+        dim_head,
+        mlp_dim,
+        dropout = 0.,
+        l1norm_after_tokens_softmax = False,
+        token_softmax_over_slots = False,
+        project_mask_groups = False,
+        gate_attention = True
+    ):
+        super().__init__()
+        self.interactions = interactions
+
+        self.attns = ModuleList([MutualAttention(
+            dim = dim,
+            num_slots = seq_lengths[j],
+            heads = heads,
+            dim_head = dim_head,
+            mlp_dim = mlp_dim,
+            dropout = dropout,
+            l1norm_after_tokens_softmax = l1norm_after_tokens_softmax,
+            token_softmax_over_slots = token_softmax_over_slots,
+            project_mask_groups = project_mask_groups,
+            gate_attention = gate_attention
+        ) for _, j in interactions])
+
+        self.norms = ModuleList([LayerNormNoBias(dim) for _ in range(num_hierarchies)])
+        self.mlps = ModuleList([FeedForward(dim, mlp_dim, dropout = dropout) for _ in range(num_hierarchies)])
+
+    def forward(self, x, masks):
+        norm_x = [norm(seq) for norm, seq in zip(self.norms, x)]
+
+        delta_x = [0.] * len(x)
+        next_masks = []
+
+        for mask, (i, j), attn in zip(masks, self.interactions, self.attns):
+            tokens_out, slots_out, next_mask = attn(norm_x[i], norm_x[j], mask)
+
+            delta_x[i] = delta_x[i] + tokens_out
+            delta_x[j] = delta_x[j] + slots_out
+            next_masks.append(next_mask)
+
+        return [seq + delta + mlp(seq + delta) for seq, delta, mlp in zip(x, delta_x, self.mlps)], next_masks
+
+# Yoshihashi et al. - https://arxiv.org/abs/2605.12021
+
+class WWT(Module):
+    def __init__(
+        self,
+        *,
+        image_size,
+        patch_size,
+        num_classes,
+        dim,
+        depth,
+        num_slots: int | tuple[int, ...],
+        sigreg_slots: bool | tuple[bool, ...] | None = None,
+        interactions: tuple[tuple[int, int], ...] | None = None,
+        heads = 8,
+        dim_head = 64,
+        mlp_dim = None,
+        channels = 3,
+        dropout = 0.,
+        return_tokens = False,
+        l1norm_after_tokens_softmax = False,
+        token_softmax_over_slots = False,
+        project_mask_groups = False,
+        gate_attention = True,
+        num_register_tokens = 0,
+        num_register_slots: int | tuple[int, ...] = 0,
+        task_heads: Module | tuple[Module, ...] | list[Module] = (),
+    ):
+        super().__init__()
+        image_height, image_width = pair(image_size)
+        patch_height, patch_width = pair(patch_size)
+
+        assert divisible_by(image_height, patch_height) and divisible_by(image_width, patch_width), 'Image dimensions must be divisible by the patch size.'
+
+        num_patches = (image_height // patch_height) * (image_width // patch_width)
+        patch_dim = channels * patch_height * patch_width
+        mlp_dim = default(mlp_dim, dim * 4)
+
+        self.to_patch_embedding = nn.Sequential(
+            Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_height, p2 = patch_width),
+            LayerNormNoBias(patch_dim),
+            nn.Linear(patch_dim, dim),
+            LayerNormNoBias(dim),
+        )
+
+        self.pos_embedding = nn.Parameter(torch.randn(num_patches, dim))
+
+        num_slots = cast_tuple(num_slots)
+
+        # ensure slots are in decreasing order to establish a part-whole hierarchy
+
+        for s1, s2 in zip(num_slots[:-1], num_slots[1:]):
+            assert s1 > s2, 'to establish a part-whole hierarchy, the number of slots must be strictly decreasing across levels'
+        num_hierarchies = 1 + len(num_slots)
+
+        if exists(sigreg_slots):
+            sigreg_slots = cast_tuple(sigreg_slots, len(num_slots))
+            assert len(sigreg_slots) == len(num_slots)
+
+        self.sigreg_slots = sigreg_slots
+
+        self.interactions = default(interactions, tuple((0, i + 1) for i in range(len(num_slots))))
+        self.interactions = tuple(tuple(interaction) for interaction in self.interactions)
+        assert len(set(self.interactions)) == len(self.interactions), 'interactions must be unique'
+
+        for i, j in self.interactions:
+            assert i < j, 'each interaction must be in strictly ascending order (from lower index to higher index)'
+        self.slots = ParameterList([nn.Parameter(torch.randn(n, dim)) for n in num_slots])
+
+        num_register_slots = cast_tuple(num_register_slots, len(num_slots))
+        assert len(num_register_slots) == len(num_slots)
+
+        self.register_tokens = nn.Parameter(torch.randn(num_register_tokens, dim))
+        self.register_slots = ParameterList([nn.Parameter(torch.randn(n, dim)) for n in num_register_slots])
+        self.num_regs = (num_register_tokens, *num_register_slots)
+
+        self.heads = heads
+        self.token_softmax_over_slots = token_softmax_over_slots
+        self.q_groups = 2 if token_softmax_over_slots else 1
+
+        self.project_mask_groups = project_mask_groups and token_softmax_over_slots
+        self.mask_groups = 1 if self.project_mask_groups else self.q_groups
+
+        self.seq_lengths = (num_patches + num_register_tokens, *(s + r for s, r in zip(num_slots, num_register_slots)))
+
+        self.layers = ModuleList([WWTBlock(
+            dim = dim,
+            num_hierarchies = num_hierarchies,
+            seq_lengths = self.seq_lengths,
+            interactions = self.interactions,
+            heads = heads,
+            dim_head = dim_head,
+            mlp_dim = mlp_dim,
+            dropout = dropout,
+            l1norm_after_tokens_softmax = l1norm_after_tokens_softmax,
+            token_softmax_over_slots = token_softmax_over_slots,
+            project_mask_groups = project_mask_groups,
+            gate_attention = gate_attention
+        ) for _ in range(depth)])
+
+        self.mlp_head = nn.Sequential(LayerNormNoBias(dim), nn.Linear(dim, num_classes))
+
+        if isinstance(task_heads, Module):
+            task_heads = (task_heads,)
+
+        self.task_heads = ModuleList(task_heads)
+        self.has_task_heads = not is_empty(self.task_heads)
+        self.return_tokens = return_tokens
+
+        if self.return_tokens:
+            self.mlp_head_tokens = nn.Sequential(LayerNormNoBias(dim), nn.Linear(dim, num_classes))
+
+    def forward(self, img, return_embeddings = False):
+        b = img.shape[0]
+
+        # get tokens
+
+        tokens = self.to_patch_embedding(img) + self.pos_embedding
+        x = [
+            tokens,
+            *(repeat(p, 's d -> b s d', b = b) for p in self.slots)
+        ]
+
+        # pack registers
+
+        all_regs = [
+            repeat(self.register_tokens, 'n d -> b n d', b = b),
+            *(repeat(p, 'n d -> b n d', b = b) for p in self.register_slots)
+        ]
+
+        packs = (pack_with_inverse([reg, seq], 'b * d') for reg, seq in zip(all_regs, x))
+        x, inverse_packs = zip(*packs)
+
+        # initial masks
+
+        masks = [tokens.new_zeros(b, self.mask_groups, self.heads, self.seq_lengths[i], self.seq_lengths[j]) for (i, j) in self.interactions]
+
+        # layers
+
+        for block in self.layers:
+            x, masks = block(x, masks)
+
+        # unpack registers
+
+        unpacked_seqs = (inv(seq) for seq, inv in zip(x, inverse_packs))
+        tokens_out, *slots_out = (seq for _, seq in unpacked_seqs)
+        slots_out = tuple(slots_out)
+        hierarchy_features = (tokens_out, *slots_out)
+
+        # process masks
+
+        processed_masks = []
+        for mask, (i, j) in zip(masks, self.interactions):
+            mask = mask[..., self.num_regs[i]:, self.num_regs[j]:]
+            if not self.token_softmax_over_slots or self.project_mask_groups:
+                mask = rearrange(mask, 'b 1 h t s -> b h t s')
+            processed_masks.append(mask)
+
+        if return_embeddings:
+            return WWTFeatureReturn(slots_out, tokens_out if self.return_tokens else None, processed_masks)
+
+        # classification
+
+        pooled_slot_logits = sum(reduce(self.mlp_head(s), 'b s c -> b c', 'mean') for s in slots_out) / len(slots_out)
+
+        if not self.return_tokens:
+            out = pooled_slot_logits
+        else:
+            pooled_token_logits = reduce(self.mlp_head_tokens(tokens_out), 'b t c -> b c', 'mean')
+            out = WWTReturn(pooled_slot_logits, pooled_token_logits)
+
+        # sigreg loss
+
+        sreg_loss = None
+        if exists(self.sigreg_slots):
+            sreg_losses = [sigreg_loss(s) for s, compute in zip(slots_out, self.sigreg_slots) if compute]
+            if not is_empty(sreg_losses):
+                sreg_loss = sum(sreg_losses)
+
+        # task heads
+
+        ret = (out,) if self.has_task_heads else out
+
+        if self.has_task_heads:
+            ret = (*ret, *(head(hierarchy_features, processed_masks, self.interactions) for head in self.task_heads))
+
+        if exists(sreg_loss):
+            ret = (*ret, sreg_loss) if isinstance(ret, tuple) else (ret, sreg_loss)
+
+        return ret
+
+if __name__ == '__main__':
+    configs = (
+        # token_softmax_over_slots, project_mask_groups, channel_first
+        (False, False, False),
+        (True, False, False),
+        (True, True, True),
+    )
+
+    for token_softmax_over_slots, project_mask_groups, channel_first in configs:
+        print(f"Testing with token_softmax_over_slots = {token_softmax_over_slots}, project_mask_groups = {project_mask_groups}, channel_first = {channel_first}")
+
+        autoencoding_head = AutoencodingHead(
+            image_size = 256,
+            patch_size = 32,
+            pathways = ((3, 2, 0), (0, 2, 3)),
+            channel_first = channel_first
+        )
+
+        model = WWT(
+            image_size = 256,
+            patch_size = 32,
+            num_classes = 1000,
+            dim = 256,
+            depth = 2,
+            num_slots = (64, 32, 16),
+            sigreg_slots = (True, False, True),
+            interactions = ((0, 1), (0, 2), (1, 2), (2, 3)),
+            heads = 4,
+            mlp_dim = 512,
+            return_tokens = True,
+            l1norm_after_tokens_softmax = True,
+            token_softmax_over_slots = token_softmax_over_slots,
+            project_mask_groups = project_mask_groups,
+            num_register_tokens = 4,
+            num_register_slots = (4, 4, 2),
+            task_heads = autoencoding_head
+        )
+
+        img = torch.randn(1, 3, 256, 256)
+
+        out, dense_feature_maps, sreg_loss = model(img)
+        slot_preds, token_preds = out
+
+        assert slot_preds.shape == (1, 1000)
+        assert token_preds.shape == (1, 1000)
+
+        assert isinstance(dense_feature_maps, tuple) and len(dense_feature_maps) == 2
+        dense_320, dense_023 = dense_feature_maps
+
+        if channel_first:
+            assert dense_320.shape == (1, 256, 8, 8)
+        else:
+            assert dense_320.shape == (1, 8, 8, 256)
+
+        assert dense_023.shape == (1, 16, 256)
+
+        assert sreg_loss.ndim == 0
+
+        print('success')

vit_pytorch/xcit.py

@@ -0,0 +1,283 @@
+from random import randrange
+
+import torch
+from torch import nn, einsum
+from torch.nn import Module, ModuleList
+import torch.nn.functional as F
+
+from einops import rearrange, repeat, pack, unpack
+from einops.layers.torch import Rearrange
+
+# helpers
+
+def exists(val):
+    return val is not None
+
+def pack_one(t, pattern):
+    return pack([t], pattern)
+
+def unpack_one(t, ps, pattern):
+    return unpack(t, ps, pattern)[0]
+
+def l2norm(t):
+    return F.normalize(t, dim = -1, p = 2)
+
+def dropout_layers(layers, dropout):
+    if dropout == 0:
+        return layers
+
+    num_layers = len(layers)
+    to_drop = torch.zeros(num_layers).uniform_(0., 1.) < dropout
+
+    # make sure at least one layer makes it
+    if all(to_drop):
+        rand_index = randrange(num_layers)
+        to_drop[rand_index] = False
+
+    layers = [layer for (layer, drop) in zip(layers, to_drop) if not drop]
+    return layers
+
+# classes
+
+class LayerScale(Module):
+    def __init__(self, dim, fn, depth):
+        super().__init__()
+        if depth <= 18:
+            init_eps = 0.1
+        elif 18 > depth <= 24:
+            init_eps = 1e-5
+        else:
+            init_eps = 1e-6
+
+        self.fn = fn
+        self.scale = nn.Parameter(torch.full((dim,), init_eps))
+
+    def forward(self, x, **kwargs):
+        return self.fn(x, **kwargs) * self.scale
+
+class FeedForward(Module):
+    def __init__(self, dim, hidden_dim, dropout = 0.):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Linear(dim, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, dim),
+            nn.Dropout(dropout)
+        )
+    def forward(self, x):
+        return self.net(x)
+
+class Attention(Module):
+    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):
+        super().__init__()
+        inner_dim = dim_head * heads
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+
+        self.norm = nn.LayerNorm(dim)
+        self.to_q = nn.Linear(dim, inner_dim, bias = False)
+        self.to_kv = nn.Linear(dim, inner_dim * 2, bias = False)
+
+        self.attend = nn.Softmax(dim = -1)
+        self.dropout = nn.Dropout(dropout)
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, dim),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x, context = None):
+        h = self.heads
+
+        x = self.norm(x)
+        context = x if not exists(context) else torch.cat((x, context), dim = 1)
+
+        qkv = (self.to_q(x), *self.to_kv(context).chunk(2, dim = -1))
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), qkv)
+
+        sim = einsum('b h i d, b h j d -> b h i j', q, k) * self.scale
+
+        attn = self.attend(sim)
+        attn = self.dropout(attn)
+
+        out = einsum('b h i j, b h j d -> b h i d', attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)
+
+class XCAttention(Module):
+    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):
+        super().__init__()
+        inner_dim = dim_head * heads
+        self.heads = heads
+        self.norm = nn.LayerNorm(dim)
+
+        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
+
+        self.temperature = nn.Parameter(torch.ones(heads, 1, 1))
+
+        self.attend = nn.Softmax(dim = -1)
+        self.dropout = nn.Dropout(dropout)
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, dim),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x):
+        h = self.heads
+        x, ps = pack_one(x, 'b * d')
+
+        x = self.norm(x)
+        q, k, v = self.to_qkv(x).chunk(3, dim = -1)
+
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h d n', h = h), (q, k, v))
+
+        q, k = map(l2norm, (q, k))
+
+        sim = einsum('b h i n, b h j n -> b h i j', q, k) * self.temperature.exp()
+
+        attn = self.attend(sim)
+        attn = self.dropout(attn)
+
+        out = einsum('b h i j, b h j n -> b h i n', attn, v)
+        out = rearrange(out, 'b h d n -> b n (h d)')
+
+        out = unpack_one(out, ps, 'b * d')
+        return self.to_out(out)
+
+class LocalPatchInteraction(Module):
+    def __init__(self, dim, kernel_size = 3):
+        super().__init__()
+        assert (kernel_size % 2) == 1
+        padding = kernel_size // 2
+
+        self.net = nn.Sequential(
+            nn.LayerNorm(dim),
+            Rearrange('b h w c -> b c h w'),
+            nn.Conv2d(dim, dim, kernel_size, padding = padding, groups = dim),
+            nn.BatchNorm2d(dim),
+            nn.GELU(),
+            nn.Conv2d(dim, dim, kernel_size, padding = padding, groups = dim),
+            Rearrange('b c h w -> b h w c'),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+class Transformer(Module):
+    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0., layer_dropout = 0.):
+        super().__init__()
+        self.layers = ModuleList([])
+        self.layer_dropout = layer_dropout
+
+        for ind in range(depth):
+            layer = ind + 1
+            self.layers.append(ModuleList([
+                LayerScale(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout), depth = layer),
+                LayerScale(dim, FeedForward(dim, mlp_dim, dropout = dropout), depth = layer)
+            ]))
+
+    def forward(self, x, context = None):
+        layers = dropout_layers(self.layers, dropout = self.layer_dropout)
+
+        for attn, ff in layers:
+            x = attn(x, context = context) + x
+            x = ff(x) + x
+
+        return x
+
+class XCATransformer(Module):
+    def __init__(self, dim, depth, heads, dim_head, mlp_dim, local_patch_kernel_size = 3, dropout = 0., layer_dropout = 0.):
+        super().__init__()
+        self.layers = ModuleList([])
+        self.layer_dropout = layer_dropout
+
+        for ind in range(depth):
+            layer = ind + 1
+            self.layers.append(ModuleList([
+                LayerScale(dim, XCAttention(dim, heads = heads, dim_head = dim_head, dropout = dropout), depth = layer),
+                LayerScale(dim, LocalPatchInteraction(dim, local_patch_kernel_size), depth = layer),
+                LayerScale(dim, FeedForward(dim, mlp_dim, dropout = dropout), depth = layer)
+            ]))
+
+    def forward(self, x):
+        layers = dropout_layers(self.layers, dropout = self.layer_dropout)
+
+        for cross_covariance_attn, local_patch_interaction, ff in layers:
+            x = cross_covariance_attn(x) + x
+            x = local_patch_interaction(x) + x
+            x = ff(x) + x
+
+        return x
+
+class XCiT(Module):
+    def __init__(
+        self,
+        *,
+        image_size,
+        patch_size,
+        num_classes,
+        dim,
+        depth,
+        cls_depth,
+        heads,
+        mlp_dim,
+        dim_head = 64,
+        dropout = 0.,
+        emb_dropout = 0.,
+        local_patch_kernel_size = 3,
+        layer_dropout = 0.
+    ):
+        super().__init__()
+        assert image_size % patch_size == 0, 'Image dimensions must be divisible by the patch size.'
+
+        num_patches = (image_size // patch_size) ** 2
+        patch_dim = 3 * patch_size ** 2
+
+        self.to_patch_embedding = nn.Sequential(
+            Rearrange('b c (h p1) (w p2) -> b h w (p1 p2 c)', p1 = patch_size, p2 = patch_size),
+            nn.LayerNorm(patch_dim),
+            nn.Linear(patch_dim, dim),
+            nn.LayerNorm(dim)
+        )
+
+        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches, dim))
+        self.cls_token = nn.Parameter(torch.randn(dim))
+
+        self.dropout = nn.Dropout(emb_dropout)
+
+        self.xcit_transformer = XCATransformer(dim, depth, heads, dim_head, mlp_dim, local_patch_kernel_size, dropout, layer_dropout)
+
+        self.final_norm = nn.LayerNorm(dim)
+
+        self.cls_transformer = Transformer(dim, cls_depth, heads, dim_head, mlp_dim, dropout, layer_dropout)
+
+        self.mlp_head = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Linear(dim, num_classes)
+        )
+
+    def forward(self, img):
+        x = self.to_patch_embedding(img)
+
+        x, ps = pack_one(x, 'b * d')
+
+        b, n, _ = x.shape
+        x += self.pos_embedding[:, :n]
+
+        x = unpack_one(x, ps, 'b * d')
+
+        x = self.dropout(x)
+
+        x = self.xcit_transformer(x)
+
+        x = self.final_norm(x)
+
+        cls_tokens = repeat(self.cls_token, 'd -> b 1 d', b = b)
+
+        x = rearrange(x, 'b ... d -> b (...) d')
+        cls_tokens = self.cls_transformer(cls_tokens, context = x)
+
+        return self.mlp_head(cls_tokens[:, 0])