canvit-nnx

The JAX and Flax NNX implementation of CanViT (Canvas Vision Transformer). Read the root backend documentation for the cross-backend API and tensor conventions.

Install

uv add canvit-nnx

Load a Hugging Face checkpoint

Load a pretrained model and an ADE20K probe from Hugging Face. Replace path/to/image.jpg with your image path.

import jax
import jax.numpy as jnp
from PIL import Image

from canvit_nnx import CanViTForSemanticSegmentation, Viewpoint, sample_at_viewpoint
from canvit_nnx.preprocess import preprocess

scene = jnp.asarray(preprocess(512)(Image.open("path/to/image.jpg").convert("RGB")))[None]
model = CanViTForSemanticSegmentation.from_pretrained_with_probe(
    pretrained_repo="canvit/canvitb16-add-vpe-pretrain-g128px-s512px-in21k-dv3b16-2026-02-02-nnx",
    probe_repo="canvit/probe-ade20k-40k-s512-c64-in21k-nnx",
)
model.eval()  # Flax NNX mutates the module and returns None.
viewpoint = Viewpoint.full_scene(batch_size=1)
state = model.init_state(batch_size=1, canvas_grid_size=64)
glimpse = sample_at_viewpoint(spatial=scene, viewpoint=viewpoint, glimpse_size_px=128)
logits, state = model(glimpse=glimpse, state=state, viewpoint=viewpoint)
jax.block_until_ready(logits)
print(logits.shape)  # [1, 64, 64, 150]

The model and probe remain separate checkpoints. They use NHWC arrays: scenes and glimpses are [B, H, W, 3], and the wrapper returns canvas logits as [B, G, G, num_classes]. Viewpoint centers use (row, col) scene coordinates in [-1, 1]; each scale is the crop's half-side.

The example uses the ImageNet-21k CanViT-B checkpoint and the ADE20K probe. The Hub lists every NNX checkpoint by its nnx tag.

Constructing a model without a checkpoint requires an explicit parameter RNG, for example:

from flax import nnx
from canvit_nnx import CanViT, CanViTConfig

model = CanViT(CanViTConfig(), rngs=nnx.Rngs(0))

Fine-tune on ImageNet-1k

canvit_nnx.specialize.in1k fine-tunes CanViT-B for ImageNet-1k classification with the paper's recipe, in bfloat16 with fused Pallas kernels on TPU. It is new in canvit-nnx 0.3.0 and needs the train extra; its README describes the data and the command.

Detect objects on COCO

canvit_nnx.detect trains CanViT-B with a CenterNet head for COCO 2017 object detection and scores average precision after each glimpse, on one host or a multi-host TPU slice. It is new in canvit-nnx 0.3.0 and needs the train extra; its README describes the data and the commands.

Develop from a checkout

cd canvit-nnx
uv sync

The checkout resolves canvit-core from its sibling directory. To convert a PyTorch checkpoint and run the conversion checks from the repository root:

uv run --project canvit-nnx python -m tools.convert_checkpoints --help