CanViT-NNX

Experimental. Port of CanViT to JAX using Flax NNX. Reference implementation: CanViT-PyTorch. May break at any time.

JAX / Flax NNX implementation of the Canvas Vision Transformer (CanViT).

Checkpoints: canvit/canvit-nnx on HuggingFace.

Install

uv add canvit-nnx

We recommend uv for dependency management.

Usage

Classification

from canvit_nnx import CanViTForImageClassification, Viewpoint, sample_at_viewpoint

clf = CanViTForImageClassification.from_pretrained(
    "canvit/canvitb16-add-vpe-finetune-g128px-s512px-in1k-2026-04-06-nnx"
)
state = clf.init_state(batch_size=1, canvas_grid_size=32)

vp = Viewpoint.full_scene(batch_size=1)
glimpse = sample_at_viewpoint(spatial=image, viewpoint=vp, glimpse_size_px=128)
logits, state = clf(glimpse, state, vp)

Pretrained backbone (dense features)

import jax.numpy as jnp
from canvit_nnx import from_pretrained, Viewpoint, sample_at_viewpoint

model = from_pretrained(
    "canvit/canvitb16-add-vpe-pretrain-g128px-s512px-in21k-dv3b16-2026-02-02-nnx"
)
state = model.init_state(batch_size=1, canvas_grid_size=32)

vp = Viewpoint.full_scene(batch_size=1)
glimpse = sample_at_viewpoint(spatial=image, viewpoint=vp, glimpse_size_px=128)
out = model(glimpse, state, vp)
state = out.state

# Canvas features should be layernormed before downstream use (PCA, probing, etc.)
canvas = model.get_spatial(state.canvas)
mean = canvas.mean(axis=-1, keepdims=True)
canvas = (canvas - mean) / jnp.sqrt(canvas.var(axis=-1, keepdims=True) + 1e-5)

Structure

canvit_nnx/
  model.py            — CanViT architecture (self-contained, no external model deps)
  classification.py   — CanViTForImageClassification
  viewpoint.py        — sample_at_viewpoint (bilinear glimpse extraction)
  hub.py              — HuggingFace Hub loading (JAX-native safetensors, no torch)

scripts/
  demo.py                  — classification demo
  verify.py                — cross-framework numerical verification
  convert_from_pytorch.py  — one-time PyTorch -> JAX checkpoint conversion
  push_to_hub.py           — convert + push to HuggingFace

Tests

uv run pytest tests/ -v

Verify against PyTorch

uv run --extra verify python scripts/verify.py
uv run --extra verify python scripts/verify.py --image path/to/image.jpg

Demo

uv run --extra demo python scripts/demo.py --image path/to/image.jpg

Citation

@article{berreby2026canvit,
  title={CanViT: Toward Active-Vision Foundation Models},
  author={Berreby, Yoha{\"i}-Eliel and Du, Sabrina and Durand, Audrey and Krishna, B. Suresh},
  year={2026},
  eprint={2603.22570},
  archivePrefix={arXiv},
  primaryClass={cs.CV}
}