Amit Moryossef a1ee1daa1a optimize NaViT with SDPA and vectorized forward pass (#353) ...

- Replace manual attention with F.scaled_dot_product_attention
- Use repeat_interleave instead of meshgrid for position computation
- Build image_ids efficiently with repeat_interleave instead of F.pad
- Remove unused Rearrange import

~56% speedup (91ms -> 58ms on 512 variable-sized images)
Numerically equivalent (max diff ~5e-4, within flash attention tolerance)

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-authored-by: Claude <noreply@anthropic.com>

2025-12-06 04:56:40 -08:00

..

__init__.py

no longer needed with einops 0.7

2023-10-19 18:16:46 -07:00

accept_video_wrapper.py

able to project the image embedding before applying time positional embedding for accept video wrapper

2025-08-13 10:15:18 -07:00

ats_vit.py

just remove PreNorm wrapper from all ViTs, as it is unlikely to change at this point

2023-08-14 09:48:55 -07:00

cait.py

just remove PreNorm wrapper from all ViTs, as it is unlikely to change at this point

2023-08-14 09:48:55 -07:00

cct_3d.py

add option to set frame padding for 3D CCT (#339)

2025-01-04 07:55:27 -08:00

cct.py

Make all CCT regularization parameters user-configurable. (#346)

2025-08-17 18:07:25 -07:00

cross_vit.py

address https://github.com/lucidrains/vit-pytorch/issues/304

2024-04-17 09:40:03 -07:00

crossformer.py

maxvit intent to build (#211)

2022-04-06 16:12:17 -07:00

cvt.py

allow channels to be customizable for cvt

2023-10-25 14:47:58 -07:00

deepvit.py

just remove PreNorm wrapper from all ViTs, as it is unlikely to change at this point

2023-08-14 09:48:55 -07:00

dino.py

implement SOTA new self-supervised learning technique from facebook for vision transformers, Dino

2021-05-02 14:00:36 -07:00

distill.py

fix t2t vit having two layernorms, and make final layernorm in distillation wrapper configurable, default to False for vit

2024-06-11 15:12:53 -07:00

efficient.py

double down on dual patch norm, fix MAE and Simmim to be compatible with dual patchnorm

2023-02-10 10:39:50 -08:00

es_vit.py

add EsViT, by popular request, an alternative to Dino that is compatible with efficient ViTs with accounting for regional self-supervised loss

2022-05-03 10:29:29 -07:00

extractor.py

offer way for extractor to return latents without detaching them

2022-07-16 16:22:40 -07:00

jumbo_vit.py

add the proposed jumbo vit from Fuller et al. of Carleton University

2025-03-05 10:50:34 -08:00

learnable_memory_vit.py

seeing a signal with dual patchnorm in another repository, fully incorporate

2023-02-06 09:45:12 -08:00

levit.py

maxvit intent to build (#211)

2022-04-06 16:12:17 -07:00

local_vit.py

just remove PreNorm wrapper from all ViTs, as it is unlikely to change at this point

2023-08-14 09:48:55 -07:00

look_vit.py

attention re-use in lookup vit should use pre-softmax attention matrix

2024-07-19 19:23:38 -07:00

mae.py

Support SimpleViT as encoder in MAE (#272)

2023-07-24 06:43:01 -07:00

max_vit_with_registers.py

give a learned bias to and from registers for maxvit + register token variant

2023-10-06 10:40:26 -07:00

max_vit.py

improvise a max vit with register tokens

2023-10-06 10:27:36 -07:00

mobile_vit.py

just remove PreNorm wrapper from all ViTs, as it is unlikely to change at this point

2023-08-14 09:48:55 -07:00

mp3.py

separate a simple vit from mp3, so that simple vit can be used after being pretrained

2023-03-07 19:31:10 -08:00

mpp.py

fix mpp

2023-06-28 08:02:43 -07:00

na_vit_nested_tensor_3d.py

address https://github.com/lucidrains/vit-pytorch/issues/352

2025-12-02 05:21:52 -08:00

na_vit_nested_tensor.py

allow for qk norm to be turned off for na vit nested tensor

2024-11-20 10:59:22 -08:00

na_vit.py

optimize NaViT with SDPA and vectorized forward pass (#353)

2025-12-06 04:56:40 -08:00

nest.py

just remove PreNorm wrapper from all ViTs, as it is unlikely to change at this point

2023-08-14 09:48:55 -07:00

normalized_vit.py

fix multiheaded qk rmsnorm in nViT

2024-10-10 19:15:17 -07:00

parallel_vit.py

just remove PreNorm wrapper from all ViTs, as it is unlikely to change at this point

2023-08-14 09:48:55 -07:00

pit.py

just remove PreNorm wrapper from all ViTs, as it is unlikely to change at this point

2023-08-14 09:48:55 -07:00

recorder.py

return None from extractor if no attention layers

2022-01-28 17:49:58 -08:00

regionvit.py

address https://github.com/lucidrains/vit-pytorch/issues/330

2024-10-04 07:01:48 -07:00

rvt.py

add a hypersphere vit, adapted from https://arxiv.org/abs/2410.01131

2024-10-09 07:32:25 -07:00

scalable_vit.py

just remove PreNorm wrapper from all ViTs, as it is unlikely to change at this point

2023-08-14 09:48:55 -07:00

sep_vit.py

just remove PreNorm wrapper from all ViTs, as it is unlikely to change at this point

2023-08-14 09:48:55 -07:00

simmim.py

double down on dual patch norm, fix MAE and Simmim to be compatible with dual patchnorm

2023-02-10 10:39:50 -08:00

simple_flash_attn_vit_3d.py

address https://github.com/lucidrains/vit-pytorch/issues/352

2025-12-02 05:21:52 -08:00

simple_flash_attn_vit.py

add a version of simple vit using flash attention

2023-03-18 09:41:39 -07:00

simple_uvit.py

add the u-vit implementation with simple vit + register tokens

2024-08-07 08:45:57 -07:00

simple_vit_1d.py

address https://github.com/lucidrains/vit-pytorch/pull/274

2023-08-09 07:53:38 -07:00

simple_vit_3d.py

address https://github.com/lucidrains/vit-pytorch/issues/352

2025-12-02 05:21:52 -08:00

simple_vit_with_fft.py

address https://github.com/lucidrains/vit-pytorch/issues/292

2023-12-23 08:11:39 -08:00

simple_vit_with_hyper_connections.py

add a simple vit flavor for a new bytedance paper that proposes to break out of the traditional one residual stream architecture - "hyper-connections"

2024-12-20 17:43:50 -08:00

simple_vit_with_patch_dropout.py

address https://github.com/lucidrains/vit-pytorch/pull/274

2023-08-09 07:53:38 -07:00

simple_vit_with_qk_norm.py

add a simple vit with qknorm, since authors seem to be promoting the technique on twitter

2023-08-14 07:58:45 -07:00

simple_vit_with_register_tokens.py

document in readme later

2023-10-03 09:26:02 -07:00

simple_vit_with_value_residual.py

make value residual learned

2024-11-24 08:21:28 -08:00

simple_vit.py

fix linear head in simple vit, thanks to @atkos

2023-08-10 14:36:21 -07:00

t2t.py

fix t2t vit having two layernorms, and make final layernorm in distillation wrapper configurable, default to False for vit

2024-06-11 15:12:53 -07:00

twins_svt.py

just remove PreNorm wrapper from all ViTs, as it is unlikely to change at this point

2023-08-14 09:48:55 -07:00

vaat.py

address https://github.com/lucidrains/vit-pytorch/issues/351

2025-11-27 06:07:43 -08:00

vat.py

register tokens for the AST in VAAT

2025-11-22 08:12:01 -08:00

vit_1d.py

Fix typo in LayerNorm (#285)

2023-10-24 12:47:21 -07:00

vit_3d.py

address https://github.com/lucidrains/vit-pytorch/issues/352

2025-12-02 05:21:52 -08:00

vit_for_small_dataset.py

just remove PreNorm wrapper from all ViTs, as it is unlikely to change at this point

2023-08-14 09:48:55 -07:00

vit_nd_rotary.py

for n-dimensional vit, have a method for fetching muon friendly parameters

2025-10-13 12:07:48 -07:00

vit_nd.py

with assistance from claude (yes it did the einops equation building here), generalize to n-dimensions

2025-09-21 06:22:43 -07:00

vit_with_decorr.py

an option to return zero for decorr aux loss if insufficient samples

2025-11-09 10:08:06 -08:00

vit_with_patch_dropout.py

just remove PreNorm wrapper from all ViTs, as it is unlikely to change at this point

2023-08-14 09:48:55 -07:00

vit_with_patch_merger.py

address https://github.com/lucidrains/vit-pytorch/pull/274

2023-08-09 07:53:38 -07:00

vit.py

fix positional embed for mean pool case and cleanup

2025-11-27 17:01:47 -08:00

vivit.py

address https://github.com/lucidrains/vit-pytorch/issues/352

2025-12-02 05:21:52 -08:00

xcit.py

add xcit (#284)

2023-10-13 09:15:13 -07:00