canvit-nnx 0.1.0__tar.gz

canvit_nnx-0.1.0/.github/workflows/release.yml

@@ -0,0 +1,22 @@
+name: Release
+
+on:
+  push:
+    tags:
+      - "v*"
+
+jobs:
+  publish:
+    runs-on: ubuntu-latest
+    environment: pypi
+    permissions:
+      id-token: write
+      contents: read
+    steps:
+      - uses: actions/checkout@v5
+
+      - uses: astral-sh/setup-uv@v5
+
+      - run: uv build
+
+      - run: uv publish

canvit_nnx-0.1.0/.gitignore

@@ -0,0 +1,20 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+
+# Virtual environments
+.venv
+
+# uv
+uv.lock
+
+# Pytest
+.pytest_cache/
+
+# Outputs
+outputs/
+main.py

canvit_nnx-0.1.0/.python-version

@@ -0,0 +1 @@
+3.12

canvit_nnx-0.1.0/CLAUDE.md

@@ -0,0 +1,50 @@
+# CLAUDE.md
+
+## Project
+
+CanViT-NNX: JAX/Flax NNX port of the Canvas Vision Transformer.
+PyPI package: `canvit-nnx`. Import: `canvit_nnx`.
+Reference implementation: [CanViT-PyTorch](https://github.com/m2b3/CanViT-PyTorch) (`canvit-pytorch` on PyPI).
+
+## Commands
+
+```bash
+uv run pytest tests/ -v                          # unit tests
+uv run --extra verify python scripts/verify.py    # cross-framework numerical check
+uv run --extra demo python scripts/demo.py --image path/to/img.jpg
+PYTHONPATH=. uv run --extra convert python scripts/push_to_hub.py  # re-push HF checkpoints
+```
+
+## Flax NNX gotchas (traps for PyTorch porters)
+
+- **LayerNorm epsilon**: Flax NNX defaults to `1e-6`, PyTorch to `1e-5`. Always explicit.
+- **GELU**: `jax.nn.gelu` defaults to `approximate=True`, PyTorch to exact. Always explicit.
+- **SDPA layout**: `jax.nn.dot_product_attention` expects BTNH, not BHTN.
+- **Linear kernel**: NNX is `(in, out)`, PyTorch is `(out, in)`. Conv: NNX `(H,W,I,O)`, PyTorch `(O,I,H,W)`.
+- **nnx.List**: plain Python lists on nnx.Module are static. Use `nnx.List([...])`.
+- **Optional modules**: `self.x = None` then `self.x = Module()` fails. Use `nnx.data(None)`.
+- **Param access**: use `param[...]` not `param.value` (deprecated in Flax 0.12+).
+- **FlatState**: iterates as `(path_tuple, leaf)` tuples, not dict. No `.items()`.
+- **Optimizer**: `nnx.Optimizer(model, tx, wrt=nnx.Param)` and `optimizer.update(model, grads)`.
+- **Pytree registration**: `@jax.tree_util.register_dataclass` for data containers passed through JIT.
+- **align_corners**: PyTorch `F.grid_sample(align_corners=False)` maps to `pixel = (grid+1)*N/2 - 0.5`.
+
+## Checkpoint loading
+
+JAX-native safetensors on HuggingFace (no torch dependency):
+- Pretrained: `canvit/canvitb16-add-vpe-pretrain-g128px-s512px-in21k-dv3b16-2026-02-02-nnx`
+- Finetuned: `canvit/canvitb16-add-vpe-finetune-g128px-s512px-in1k-2026-04-06-nnx`
+
+PyTorch-to-JAX conversion is in `scripts/convert_from_pytorch.py` (not part of the installed package).
+The PyTorch pretrained checkpoint key prefix is `backbone.` (the actual ViTBackbone).
+The PyTorch finetuned checkpoint wraps CanViT as `canvit.` — strip that, and inner keys match.
+
+## Performance notes (as of 2026-04-07, may be stale)
+
+- On CPU (M4 Pro), JAX forward was ~68ms vs torch.compile ~41ms. Initial investigation
+  suggested canvas cross-attention SDPA as the bottleneck (XLA CPU appeared to decompose
+  it into separate ops while PyTorch used a fused kernel), but this was not thoroughly
+  verified on GPU. Re-benchmark before drawing conclusions.
+- Cross-framework backward gradient differences were observed in early blocks (0-3).
+  Hypothesis was f32 matmul accumulation order differences, but not definitively proven.
+  Run `scripts/verify.py` for current state.

canvit_nnx-0.1.0/PKG-INFO

@@ -0,0 +1,133 @@
+Metadata-Version: 2.4
+Name: canvit-nnx
+Version: 0.1.0
+Summary: CanViT (Canvas Vision Transformer) -- JAX / Flax NNX
+Author-email: Yohaï-Eliel Berreby <me@yberreby.com>
+License-Expression: MIT
+Requires-Python: >=3.12
+Requires-Dist: flax>=0.12
+Requires-Dist: huggingface-hub
+Requires-Dist: jax>=0.9
+Requires-Dist: jaxlib>=0.9
+Requires-Dist: safetensors
+Provides-Extra: convert
+Requires-Dist: canvit-pytorch>=0.1.9; extra == 'convert'
+Requires-Dist: torch>=2.0; extra == 'convert'
+Provides-Extra: demo
+Requires-Dist: pillow; extra == 'demo'
+Provides-Extra: dev
+Requires-Dist: canvit-pytorch>=0.1.9; extra == 'dev'
+Requires-Dist: ipython; extra == 'dev'
+Requires-Dist: pillow; extra == 'dev'
+Requires-Dist: pypatree; extra == 'dev'
+Requires-Dist: pytest; extra == 'dev'
+Requires-Dist: torch>=2.0; extra == 'dev'
+Requires-Dist: torchvision; extra == 'dev'
+Provides-Extra: verify
+Requires-Dist: canvit-pytorch>=0.1.9; extra == 'verify'
+Requires-Dist: pillow; extra == 'verify'
+Requires-Dist: torch>=2.0; extra == 'verify'
+Description-Content-Type: text/markdown
+
+# CanViT-NNX
+
+> **Experimental.** Port of [CanViT](https://arxiv.org/abs/2603.22570) to [JAX](https://github.com/jax-ml/jax) using [Flax NNX](https://flax.readthedocs.io/). Reference implementation: [CanViT-PyTorch](https://github.com/m2b3/CanViT-PyTorch). May break at any time.
+
+[JAX](https://github.com/jax-ml/jax) / [Flax NNX](https://flax.readthedocs.io/) implementation of the [Canvas Vision Transformer (CanViT)](https://arxiv.org/abs/2603.22570).
+
+Checkpoints: [canvit/canvit-nnx](https://huggingface.co/collections/canvit/canvit-nnx-69d4864ff6ee17d87e7b91d4) on HuggingFace.
+
+## Install
+
+```bash
+uv add canvit-nnx
+```
+
+We recommend [`uv`](https://docs.astral.sh/uv/) for dependency management.
+
+## Usage
+
+### Classification
+
+```python
+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)
+
+```python
+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
+
+```bash
+uv run pytest tests/ -v
+```
+
+## Verify against PyTorch
+
+```bash
+uv run --extra verify python scripts/verify.py
+uv run --extra verify python scripts/verify.py --image path/to/image.jpg
+```
+
+## Demo
+
+```bash
+uv run --extra demo python scripts/demo.py --image path/to/image.jpg
+```
+
+## Citation
+
+```bibtex
+@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}
+}
+```

canvit_nnx-0.1.0/README.md

@@ -0,0 +1,102 @@
+# CanViT-NNX
+
+> **Experimental.** Port of [CanViT](https://arxiv.org/abs/2603.22570) to [JAX](https://github.com/jax-ml/jax) using [Flax NNX](https://flax.readthedocs.io/). Reference implementation: [CanViT-PyTorch](https://github.com/m2b3/CanViT-PyTorch). May break at any time.
+
+[JAX](https://github.com/jax-ml/jax) / [Flax NNX](https://flax.readthedocs.io/) implementation of the [Canvas Vision Transformer (CanViT)](https://arxiv.org/abs/2603.22570).
+
+Checkpoints: [canvit/canvit-nnx](https://huggingface.co/collections/canvit/canvit-nnx-69d4864ff6ee17d87e7b91d4) on HuggingFace.
+
+## Install
+
+```bash
+uv add canvit-nnx
+```
+
+We recommend [`uv`](https://docs.astral.sh/uv/) for dependency management.
+
+## Usage
+
+### Classification
+
+```python
+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)
+
+```python
+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
+
+```bash
+uv run pytest tests/ -v
+```
+
+## Verify against PyTorch
+
+```bash
+uv run --extra verify python scripts/verify.py
+uv run --extra verify python scripts/verify.py --image path/to/image.jpg
+```
+
+## Demo
+
+```bash
+uv run --extra demo python scripts/demo.py --image path/to/image.jpg
+```
+
+## Citation
+
+```bibtex
+@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}
+}
+```

canvit_nnx-0.1.0/canvit_nnx/__init__.py

@@ -0,0 +1,13 @@
+"""CanViT: minimal JAX/Flax NNX implementation."""
+
+from canvit_nnx.model import (
+    CANVIT_B16,
+    CanViT,
+    CanViTConfig,
+    CanViTOutput,
+    RecurrentState,
+    Viewpoint,
+)
+from canvit_nnx.classification import CanViTForImageClassification
+from canvit_nnx.viewpoint import sample_at_viewpoint
+from canvit_nnx.hub import from_pretrained, save_state, load_state

canvit_nnx-0.1.0/canvit_nnx/classification.py

@@ -0,0 +1,86 @@
+"""CanViT for image classification: CanViT backbone + LN -> Linear head.
+
+Mirrors the canvit_pytorch.model.classification API.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+
+import flax.nnx as nnx
+import jax
+import jax.numpy as jnp
+
+from canvit_nnx.model import CanViT, CanViTConfig, RecurrentState, Viewpoint, CanViTOutput
+from canvit_nnx.hub import load_state, FINETUNED_REPO
+
+log = logging.getLogger(__name__)
+
+Array = jax.Array
+
+
+class CanViTForImageClassification(nnx.Module):
+    """CanViT backbone + LN -> Linear classification head.
+
+    Usage::
+
+        clf = CanViTForImageClassification.from_pretrained()
+        state = clf.init_state(batch_size=1, canvas_grid_size=32)
+        logits, state = clf(glimpse, state, viewpoint)
+    """
+
+    def __init__(self, cfg: CanViTConfig, n_classes: int, *, rngs: nnx.Rngs):
+        self.canvit = CanViT(cfg, rngs=rngs)
+        self.norm = nnx.LayerNorm(cfg.embed_dim, epsilon=1e-5, rngs=rngs)
+        self.head = nnx.Linear(cfg.embed_dim, n_classes, rngs=rngs)
+
+    @property
+    def n_classes(self) -> int:
+        return self.head.out_features
+
+    @property
+    def cfg(self) -> CanViTConfig:
+        return self.canvit.cfg
+
+    def init_state(self, batch_size: int, canvas_grid_size: int) -> RecurrentState:
+        return self.canvit.init_state(batch_size, canvas_grid_size)
+
+    def __call__(self, glimpse: Array, state: RecurrentState,
+                 viewpoint: Viewpoint) -> tuple[Array, RecurrentState]:
+        """Returns (logits [B, n_classes], new_state)."""
+        out = self.canvit(glimpse, state, viewpoint)
+        cls = out.state.recurrent_cls[:, 0].astype(jnp.float32)
+        logits = self.head(self.norm(cls))
+        return logits, out.state
+
+    def canvit_forward(self, glimpse: Array, state: RecurrentState,
+                       viewpoint: Viewpoint) -> CanViTOutput:
+        """Run CanViT backbone only (for training with separate head_forward)."""
+        return self.canvit(glimpse, state, viewpoint)
+
+    def head_forward(self, cls: Array) -> Array:
+        """LN -> Linear on [B, D] CLS token."""
+        return self.head(self.norm(cls.astype(jnp.float32)))
+
+    @staticmethod
+    def from_pretrained(repo_id: str = FINETUNED_REPO) -> CanViTForImageClassification:
+        """Load finetuned CanViT classifier from HuggingFace Hub.
+
+        Uses JAX-native safetensors -- no PyTorch dependency.
+        """
+        from huggingface_hub import hf_hub_download
+
+        config_path = Path(hf_hub_download(repo_id, "config.json"))
+        weights_path = Path(hf_hub_download(repo_id, "model.safetensors"))
+
+        config = json.loads(config_path.read_text())
+        mc = config["model_config"]
+        valid = set(CanViTConfig.__dataclass_fields__)
+        cfg = CanViTConfig(**{k: v for k, v in mc.items() if k in valid})
+        n_classes = config["n_classes"]
+
+        clf = CanViTForImageClassification(cfg, n_classes, rngs=nnx.Rngs(0))
+        load_state(clf, weights_path)
+        return clf

canvit_nnx-0.1.0/canvit_nnx/hub.py

@@ -0,0 +1,92 @@
+"""HuggingFace Hub integration: load and save JAX-native CanViT checkpoints.
+
+No PyTorch dependency — loads directly from safetensors.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+
+import flax.nnx as nnx
+import jax.numpy as jnp
+import numpy as np
+
+from canvit_nnx.model import CanViT, CanViTConfig
+
+log = logging.getLogger(__name__)
+
+PRETRAINED_REPO = "canvit/canvitb16-add-vpe-pretrain-g128px-s512px-in21k-dv3b16-2026-02-02-nnx"
+FINETUNED_REPO = "canvit/canvitb16-add-vpe-finetune-g128px-s512px-in1k-2026-04-06-nnx"
+
+def save_state(module: nnx.Module, path: str | Path) -> int:
+    """Save all NNX module state to safetensors. Returns number of tensors saved."""
+    from safetensors.numpy import save_file
+
+    _, state = nnx.split(module)
+    params = {}
+    for path_tuple, leaf in state.flat_state():
+        key = ".".join(str(p) for p in path_tuple)
+        params[key] = np.array(leaf[...])
+
+    save_file(params, str(path))
+    log.info("Saved %d tensors to %s", len(params), path)
+    return len(params)
+
+
+def load_state(module: nnx.Module, path: str | Path) -> int:
+    """Load safetensors weights into an NNX module. Returns number loaded.
+
+    Every key in the checkpoint must exist in the module (and vice versa).
+    """
+    from safetensors import safe_open
+
+    _, state = nnx.split(module)
+
+    state_map: dict[str, tuple] = {}
+    for path_tuple, leaf in state.flat_state():
+        key = ".".join(str(p) for p in path_tuple)
+        state_map[key] = leaf
+
+    loaded = 0
+    with safe_open(str(path), framework="numpy") as f:
+        ckpt_keys = set(f.keys())
+        model_keys = set(state_map.keys())
+
+        missing_in_model = ckpt_keys - model_keys
+        missing_in_ckpt = model_keys - ckpt_keys
+        if missing_in_model:
+            raise ValueError(f"Checkpoint has keys not in model: {sorted(missing_in_model)}")
+        if missing_in_ckpt:
+            raise ValueError(f"Model has keys not in checkpoint: {sorted(missing_in_ckpt)}")
+
+        for key in sorted(f.keys()):
+            leaf = state_map[key]
+            arr = jnp.array(f.get_tensor(key), dtype=jnp.float32)
+            if leaf[...].shape != arr.shape:
+                raise ValueError(
+                    f"Shape mismatch at {key}: model={leaf[...].shape}, checkpoint={arr.shape}"
+                )
+            leaf[...] = arr
+            loaded += 1
+
+    log.info("Loaded %d tensors from %s", loaded, path)
+    return loaded
+
+
+def from_pretrained(repo_id: str = PRETRAINED_REPO) -> CanViT:
+    """Load pretrained CanViT from HuggingFace Hub (JAX-native checkpoint)."""
+    from huggingface_hub import hf_hub_download
+
+    config_path = hf_hub_download(repo_id, "config.json")
+    weights_path = hf_hub_download(repo_id, "model.safetensors")
+
+    config = json.loads(Path(config_path).read_text())
+    mc = config["model_config"]
+    valid = set(CanViTConfig.__dataclass_fields__)
+    cfg = CanViTConfig(**{k: v for k, v in mc.items() if k in valid})
+
+    model = CanViT(cfg, rngs=nnx.Rngs(0))
+    load_state(model, weights_path)
+    return model

canvit_nnx-0.1.0/canvit_nnx/model.py

@@ -0,0 +1,507 @@
+"""CanViT: dual-stream vision transformer with canvas cross-attention.
+
+Self-contained Flax NNX implementation. No external model dependencies.
+Reference: CanViT-PyTorch (authoritative) and CanViT-MLX (compact port).
+
+Conventions:
+  - Images are NHWC (channels-last), matching JAX/XLA native layout.
+  - Coordinates are (y, x) order in [-1, 1]², matching tensor [H, W] indexing.
+  - RoPE is always computed and applied in float32.
+  - Canvas layout: [registers | spatial].
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import NamedTuple
+
+import flax.nnx as nnx
+import jax
+import jax.numpy as jnp
+
+Array = jax.Array
+
+
+# =============================================================================
+# Config
+# =============================================================================
+
+
+@dataclass(frozen=True)
+class CanViTConfig:
+    embed_dim: int = 768
+    num_heads: int = 12
+    n_blocks: int = 12
+    patch_size: int = 16
+    ffn_ratio: float = 4.0
+    rope_base: float = 100.0
+    layerscale_init: float = 1e-5
+
+    rw_stride: int = 2
+    enable_reads: bool = True
+    n_backbone_registers: int = 5
+    n_canvas_registers: int = 16
+    canvas_num_heads: int = 8
+    canvas_head_dim: int = 128
+    enable_vpe: bool = True
+
+    @property
+    def canvas_dim(self) -> int:
+        return self.canvas_num_heads * self.canvas_head_dim
+
+    @property
+    def head_dim(self) -> int:
+        return self.embed_dim // self.num_heads
+
+
+CANVIT_B16 = CanViTConfig()
+
+
+# =============================================================================
+# Data types
+# =============================================================================
+
+
+@jax.tree_util.register_dataclass
+@dataclass
+class Viewpoint:
+    centers: Array  # [B, 2] in [-1, 1], (y, x) order, float32
+    scales: Array   # [B], float32
+
+    @staticmethod
+    def full_scene(batch_size: int) -> Viewpoint:
+        return Viewpoint(
+            centers=jnp.zeros((batch_size, 2), dtype=jnp.float32),
+            scales=jnp.ones((batch_size,), dtype=jnp.float32),
+        )
+
+
+@jax.tree_util.register_dataclass
+@dataclass
+class RecurrentState:
+    canvas: Array       # [B, n_canvas_registers + G², canvas_dim]
+    recurrent_cls: Array  # [B, 1, local_dim]
+
+
+class CanViTOutput(NamedTuple):
+    state: RecurrentState
+    local_patches: Array  # [B, H*W, local_dim]
+
+
+# =============================================================================
+# Coordinates
+# =============================================================================
+
+
+def grid_coords(H: int, W: int) -> Array:
+    """Normalized [-1, 1] cell-center coordinates. Returns [H*W, 2] in (y, x) order."""
+    y = (jnp.arange(H, dtype=jnp.float32) + 0.5) / H * 2 - 1
+    x = (jnp.arange(W, dtype=jnp.float32) + 0.5) / W * 2 - 1
+    yy, xx = jnp.meshgrid(y, x, indexing="ij")
+    return jnp.stack([yy, xx], axis=-1).reshape(H * W, 2)
+
+
+def glimpse_positions(center: Array, scale: Array, H: int, W: int) -> Array:
+    """Scene-space positions for a glimpse grid. Returns [B, H*W, 2]."""
+    retino = grid_coords(H, W)  # [H*W, 2]
+    return center[:, None, :] + scale[:, None, None] * retino[None, :, :]
+
+
+# =============================================================================
+# 2D Rotary Position Embeddings
+# =============================================================================
+
+
+def _rope_periods(head_dim: int, base: float) -> Array:
+    """Geometric frequency periods for 2D RoPE. Returns [head_dim // 4]."""
+    assert head_dim % 4 == 0
+    n = head_dim // 4
+    return base ** (jnp.arange(n, dtype=jnp.float32) / n)
+
+
+class RoPE(NamedTuple):
+    sin: Array  # [B, N, 1, head_dim]  (BTNH layout for JAX SDPA)
+    cos: Array  # [B, N, 1, head_dim]
+
+
+def compute_rope(positions: Array, head_dim: int, base: float) -> RoPE:
+    """Compute 2D RoPE sin/cos from positions [B, N, 2]. Always float32."""
+    periods = _rope_periods(head_dim, base)
+    # [B, N, 2, 1] / [n_freqs] -> [B, N, 2, n_freqs]
+    angles = 2 * math.pi * positions[..., None].astype(jnp.float32) / periods
+    B, N = angles.shape[:2]
+    angles = jnp.tile(angles.reshape(B, N, -1), (1, 1, 2))  # [B, N, head_dim]
+    sin = jnp.sin(angles)[:, :, None, :]  # [B, N, 1, head_dim]
+    cos = jnp.cos(angles)[:, :, None, :]
+    return RoPE(sin, cos)
+
+
+def apply_rope(x: Array, rope: RoPE) -> Array:
+    """Apply RoPE rotation to spatial tokens, leaving prefix tokens unchanged.
+
+    x: [B, N_total, H, head_dim] (BTNH layout). RoPE covers N_spatial tokens.
+    Rotation is in float32; result cast back to x.dtype.
+    """
+    n_prefix = x.shape[1] - rope.sin.shape[1]
+    half = x.shape[3] // 2
+    x_s = x[:, n_prefix:].astype(jnp.float32)
+    x_rot = jnp.concatenate([-x_s[..., half:], x_s[..., :half]], axis=-1)
+    spatial = (x_s * rope.cos + x_rot * rope.sin).astype(x.dtype)
+    if n_prefix > 0:
+        return jnp.concatenate([x[:, :n_prefix], spatial], axis=1)
+    return spatial
+
+
+# =============================================================================
+# Building blocks
+# =============================================================================
+
+
+class PatchEmbed(nnx.Module):
+    def __init__(self, patch_size: int, embed_dim: int, rngs: nnx.Rngs):
+        self.proj = nnx.Conv(3, embed_dim, kernel_size=(patch_size, patch_size),
+                             strides=(patch_size, patch_size), padding="VALID",
+                             use_bias=True, rngs=rngs)
+
+    def __call__(self, x: Array) -> tuple[Array, int, int]:
+        """x: [B, H, W, 3] -> patches [B, H'*W', D], H', W'."""
+        x = self.proj(x)
+        B, H, W, D = x.shape
+        return x.reshape(B, H * W, D), H, W
+
+
+class LayerScale(nnx.Module):
+    def __init__(self, dim: int, init_value: float = 1e-5):
+        self.gamma = nnx.Param(jnp.full((dim,), init_value))
+
+    def __call__(self, x: Array) -> Array:
+        return x * self.gamma
+
+
+class MLP(nnx.Module):
+    def __init__(self, dim: int, hidden_dim: int, rngs: nnx.Rngs):
+        self.fc1 = nnx.Linear(dim, hidden_dim, rngs=rngs)
+        self.fc2 = nnx.Linear(hidden_dim, dim, rngs=rngs)
+
+    def __call__(self, x: Array) -> Array:
+        # approximate=False: JAX defaults to tanh approx, PyTorch to exact
+        return self.fc2(jax.nn.gelu(self.fc1(x), approximate=False))
+
+
+class SelfAttention(nnx.Module):
+    """Multi-head self-attention with fused QKV and 2D RoPE.
+
+    Matches DINOv3 convention: K bias is zeroed out (Q and V biases active).
+    """
+
+    def __init__(self, dim: int, num_heads: int, rngs: nnx.Rngs):
+        assert dim % num_heads == 0
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.scale = self.head_dim ** -0.5
+        self.qkv = nnx.Linear(dim, dim * 3, rngs=rngs)
+        # Mask: Q bias ON, K bias OFF, V bias ON (DINOv3 convention)
+        self._bias_mask = nnx.Variable(
+            jnp.concatenate([jnp.ones(dim), jnp.zeros(dim), jnp.ones(dim)])
+        )
+        self.proj = nnx.Linear(dim, dim, rngs=rngs)
+
+    def __call__(self, x: Array, rope: RoPE) -> Array:
+        B, N, D = x.shape
+        # Manual matmul with masked bias (K bias zeroed, DINOv3 convention)
+        qkv = x @ self.qkv.kernel[...]  # NNX kernel is (in, out)
+        if self.qkv.bias is not None:
+            qkv = qkv + self.qkv.bias[...] * self._bias_mask[...]
+        qkv = qkv.reshape(B, N, 3, self.num_heads, self.head_dim)
+        q = qkv[:, :, 0]  # [B, N, H, D_h] — BTNH for JAX SDPA
+        k = qkv[:, :, 1]
+        v = qkv[:, :, 2]
+        q = apply_rope(q, rope)
+        k = apply_rope(k, rope)
+        out = jax.nn.dot_product_attention(q, k, v, scale=self.scale)
+        return self.proj(out.reshape(B, N, D))
+
+
+class ViTBlock(nnx.Module):
+    def __init__(self, dim: int, num_heads: int, ffn_ratio: float,
+                 layerscale_init: float, rngs: nnx.Rngs):
+        self.norm1 = nnx.LayerNorm(dim, epsilon=1e-5, rngs=rngs)
+        self.attn = SelfAttention(dim, num_heads, rngs=rngs)
+        self.ls1 = LayerScale(dim, layerscale_init)
+        self.norm2 = nnx.LayerNorm(dim, epsilon=1e-5, rngs=rngs)
+        self.mlp = MLP(dim, int(dim * ffn_ratio), rngs=rngs)
+        self.ls2 = LayerScale(dim, layerscale_init)
+
+    def __call__(self, x: Array, rope: RoPE) -> Array:
+        x = x + self.ls1(self.attn(self.norm1(x), rope))
+        return x + self.ls2(self.mlp(self.norm2(x)))
+
+
+# =============================================================================
+# Canvas attention (asymmetric cross-attention)
+# =============================================================================
+
+
+def _to_multihead(x: Array, num_heads: int) -> Array:
+    """[B, N, D] -> [B, N, H, head_dim] (JAX SDPA expects BTNH)."""
+    B, N, D = x.shape
+    return x.reshape(B, N, num_heads, D // num_heads)
+
+
+def _from_multihead(x: Array) -> Array:
+    """[B, N, H, head_dim] -> [B, N, D]."""
+    B, N, H, hd = x.shape
+    return x.reshape(B, N, H * hd)
+
+
+class CanvasReadAttention(nnx.Module):
+    """Local queries canvas. Dense projections on local side only (asymmetric)."""
+
+    def __init__(self, local_dim: int, canvas_dim: int, num_heads: int, rngs: nnx.Rngs):
+        assert canvas_dim % num_heads == 0
+        self.num_heads = num_heads
+        self.scale = (canvas_dim // num_heads) ** -0.5
+        self.q_norm = nnx.LayerNorm(local_dim, epsilon=1e-5, rngs=rngs)
+        self.kv_norm = nnx.LayerNorm(canvas_dim, epsilon=1e-5, rngs=rngs)
+        self.q_proj = nnx.Linear(local_dim, canvas_dim, rngs=rngs)
+        self.out_proj = nnx.Linear(canvas_dim, local_dim, rngs=rngs)
+
+    def __call__(self, query: Array, kv: Array,
+                 q_rope: RoPE, kv_rope: RoPE) -> Array:
+        q = _to_multihead(self.q_proj(self.q_norm(query)), self.num_heads)
+        kv_n = self.kv_norm(kv)
+        k = _to_multihead(kv_n, self.num_heads)  # Identity projection
+        v = _to_multihead(kv_n, self.num_heads)
+        q = apply_rope(q, q_rope)
+        k = apply_rope(k, kv_rope)
+        out = jax.nn.dot_product_attention(q, k.astype(q.dtype), v.astype(q.dtype),
+                                           scale=self.scale)
+        return self.out_proj(_from_multihead(out))
+
+
+class CanvasWriteAttention(nnx.Module):
+    """Canvas queries local. Dense projections on local side only (asymmetric).
+
+    When gate_bias_init is set, uses convex update: lerp(canvas, attn_out, gate).
+    """
+
+    def __init__(self, local_dim: int, canvas_dim: int, num_heads: int,
+                 rngs: nnx.Rngs, gate_bias_init: float | None = None):
+        assert canvas_dim % num_heads == 0
+        self.num_heads = num_heads
+        self.scale = (canvas_dim // num_heads) ** -0.5
+        self.q_norm = nnx.LayerNorm(canvas_dim, epsilon=1e-5, rngs=rngs)
+        self.kv_norm = nnx.LayerNorm(local_dim, epsilon=1e-5, rngs=rngs)
+        self.k_proj = nnx.Linear(local_dim, canvas_dim, rngs=rngs)
+        self.v_proj = nnx.Linear(local_dim, canvas_dim, rngs=rngs)
+        self.gate_linear: nnx.Linear | None = None
+        if gate_bias_init is not None:
+            self.gate_linear = nnx.Linear(canvas_dim, 1, rngs=rngs)
+            self.gate_linear.bias = nnx.Param(jnp.full((1,), gate_bias_init))
+
+    @property
+    def is_convex(self) -> bool:
+        return self.gate_linear is not None
+
+    def __call__(self, query: Array, kv: Array,
+                 q_rope: RoPE, kv_rope: RoPE) -> Array:
+        q = _to_multihead(self.q_norm(query), self.num_heads)  # Identity projection
+        kv_n = self.kv_norm(kv)
+        k = _to_multihead(self.k_proj(kv_n), self.num_heads)
+        v = _to_multihead(self.v_proj(kv_n), self.num_heads)
+        q = apply_rope(q, q_rope)
+        k = apply_rope(k, kv_rope)
+        attn_out = _from_multihead(
+            jax.nn.dot_product_attention(q, k.astype(q.dtype), v.astype(q.dtype),
+                                         scale=self.scale)
+        )
+        if self.gate_linear is None:
+            return attn_out
+        gate = jax.nn.sigmoid(self.gate_linear(attn_out))
+        return (1 - gate) * query + gate * attn_out
+
+
+# =============================================================================
+# VPE: Viewpoint Positional Encoding
+# =============================================================================
+
+
+class VPEEncoder(nnx.Module):
+    """Encodes (y, x, scale) viewpoint into a high-dimensional vector via Random Fourier Features."""
+
+    def __init__(self, rff_dim: int, *, rngs: nnx.Rngs, seed: int = 42):
+        assert rff_dim > 0 and rff_dim % 2 == 0
+        self.rff_dim = rff_dim
+        # Frozen RFF projection matrix — deterministic from seed, not from rngs
+        key = jax.random.PRNGKey(seed)
+        self.B_mat = nnx.Variable(jax.random.normal(key, (rff_dim // 2, 3), dtype=jnp.float32))
+        self.norm = nnx.LayerNorm(rff_dim, epsilon=1e-5, rngs=rngs)
+        # Match PyTorch init: scale = 1/sqrt(rff_dim)
+        self.norm.scale = nnx.Param(jnp.full((rff_dim,), 1.0 / math.sqrt(rff_dim)))
+
+    def __call__(self, y: Array, x: Array, s: Array) -> Array:
+        """y, x: [B], s: [B] -> [B, rff_dim]. Always float32."""
+        z = jnp.stack([y / s, x / s, jnp.log(s)], axis=-1).astype(jnp.float32)
+        proj = z @ self.B_mat[...].T
+        enc = jnp.concatenate([jnp.cos(proj), jnp.sin(proj)], axis=-1)
+        return self.norm(enc)
+
+
+# =============================================================================
+# CanViT: the full model
+# =============================================================================
+
+
+def _compute_rw_positions(n_blocks: int, rw_stride: int,
+                          enable_reads: bool) -> tuple[list[int], list[int]]:
+    """Schedule read/write operations across blocks.
+
+    Pattern: [rw_stride blocks] R [rw_stride blocks] W ... Always ends with W.
+    """
+    read_after: list[int] = []
+    write_after: list[int] = []
+    for i, pos in enumerate(range(rw_stride - 1, n_blocks, rw_stride)):
+        if i % 2 == 0:
+            read_after.append(pos)
+        else:
+            write_after.append(pos)
+    if not write_after or write_after[-1] != n_blocks - 1:
+        write_after.append(n_blocks - 1)
+    if not enable_reads:
+        read_after = []
+    return read_after, write_after
+
+
+class CanViT(nnx.Module):
+    """Dual-stream vision transformer with canvas cross-attention.
+
+    Processes NHWC glimpses. Canvas and CLS token persist across timesteps.
+    Generalizes to any canvas and glimpse grid sizes at runtime via RoPE.
+    """
+
+    def __init__(self, cfg: CanViTConfig = CANVIT_B16, *, rngs: nnx.Rngs):
+        self.cfg = cfg
+        d = cfg.embed_dim
+        cd = cfg.canvas_dim
+
+        # Backbone
+        self.patch_embed = PatchEmbed(cfg.patch_size, d, rngs)
+        self.blocks = nnx.List([
+            ViTBlock(d, cfg.num_heads, cfg.ffn_ratio, cfg.layerscale_init, rngs)
+            for _ in range(cfg.n_blocks)
+        ])
+
+        # Read/write schedule
+        read_after, write_after = _compute_rw_positions(
+            cfg.n_blocks, cfg.rw_stride, cfg.enable_reads)
+        self._read_after = tuple(read_after)
+        self._write_after = tuple(write_after)
+
+        # Canvas attention layers
+        self.canvas_read = nnx.List([
+            CanvasReadAttention(d, cd, cfg.canvas_num_heads, rngs)
+            for _ in read_after
+        ])
+        self.canvas_write = nnx.List([
+            CanvasWriteAttention(d, cd, cfg.canvas_num_heads, rngs)
+            for _ in write_after
+        ])
+
+        # Learnable inits
+        canvas_scale = 1.0 / math.sqrt(cd)
+        local_scale = 1.0 / math.sqrt(d)
+        self.canvas_register_init = nnx.Param(
+            jax.random.normal(rngs.params(), (1, cfg.n_canvas_registers, cd)) * canvas_scale
+        )
+        self.canvas_spatial_init = nnx.Param(
+            jax.random.normal(rngs.params(), (1, 1, cd)) * canvas_scale
+        )
+        self.recurrent_cls_init = nnx.Param(
+            jax.random.normal(rngs.params(), (1, 1, d)) * local_scale
+        )
+        self.backbone_registers = nnx.Param(
+            jax.random.normal(rngs.params(), (1, cfg.n_backbone_registers, d)) * 0.02
+        )
+
+        # VPE
+        if cfg.enable_vpe:
+            self.vpe = VPEEncoder(rff_dim=d, rngs=rngs)
+        else:
+            self.vpe = nnx.data(None)
+
+    def init_state(self, batch_size: int, canvas_grid_size: int) -> RecurrentState:
+        cfg = self.cfg
+        n_spatial = canvas_grid_size ** 2
+        regs = jnp.broadcast_to(self.canvas_register_init[...], (batch_size, cfg.n_canvas_registers, cfg.canvas_dim))
+        spatial = jnp.broadcast_to(self.canvas_spatial_init[...], (batch_size, n_spatial, cfg.canvas_dim))
+        return RecurrentState(
+            canvas=jnp.concatenate([regs, spatial], axis=1),
+            recurrent_cls=jnp.broadcast_to(self.recurrent_cls_init[...], (batch_size, 1, cfg.embed_dim)),
+        )
+
+    def get_spatial(self, canvas: Array) -> Array:
+        return canvas[:, self.cfg.n_canvas_registers:]
+
+    def __call__(self, glimpse: Array, state: RecurrentState,
+                 viewpoint: Viewpoint) -> CanViTOutput:
+        """Process one glimpse. glimpse: [B, H, W, 3] NHWC."""
+        B = glimpse.shape[0]
+        cfg = self.cfg
+        canvas = state.canvas
+
+        # Patch embed
+        patches, H, W = self.patch_embed(glimpse)
+        assert H == W
+        n_patches = H * W
+
+        # Assemble local tokens: [vpe?, cls, registers, patches]
+        parts: list[Array] = []
+        has_vpe = self.vpe is not None
+        if has_vpe:
+            vpe_tok = self.vpe(
+                viewpoint.centers[:, 0], viewpoint.centers[:, 1], viewpoint.scales
+            )
+            parts.append(vpe_tok[:, None, :].astype(patches.dtype))
+        parts.extend([
+            state.recurrent_cls,
+            jnp.broadcast_to(self.backbone_registers[...], (B, cfg.n_backbone_registers, cfg.embed_dim)),
+            patches,
+        ])
+        local = jnp.concatenate(parts, axis=1)
+
+        # RoPE for local and canvas positions
+        local_pos = glimpse_positions(viewpoint.centers, viewpoint.scales, H, W)
+        bb_rope = compute_rope(local_pos, cfg.head_dim, cfg.rope_base)
+        ca_local_rope = compute_rope(local_pos, cfg.canvas_head_dim, cfg.rope_base)
+
+        n_spatial = canvas.shape[1] - cfg.n_canvas_registers
+        cg = int(math.sqrt(n_spatial))
+        assert cg * cg == n_spatial
+        canvas_pos = jnp.broadcast_to(grid_coords(cg, cg)[None], (B, n_spatial, 2))
+        ca_canvas_rope = compute_rope(canvas_pos, cfg.canvas_head_dim, cfg.rope_base)
+
+        # Process through blocks with interleaved read/write
+        ri, wi = 0, 0
+        for bi in range(cfg.n_blocks):
+            local = self.blocks[bi](local, bb_rope)
+
+            if ri < len(self._read_after) and bi == self._read_after[ri]:
+                local = local + self.canvas_read[ri](local, canvas, ca_local_rope, ca_canvas_rope)
+                ri += 1
+
+            if wi < len(self._write_after) and bi == self._write_after[wi]:
+                write_out = self.canvas_write[wi](canvas, local, ca_canvas_rope, ca_local_rope)
+                canvas = write_out if self.canvas_write[wi].is_convex else canvas + write_out
+                wi += 1
+
+        # Unpack
+        idx = (1 if has_vpe else 0)
+        new_cls = local[:, idx:idx + 1]
+        idx += 1 + cfg.n_backbone_registers
+        out_patches = local[:, idx:idx + n_patches]
+
+        return CanViTOutput(
+            state=RecurrentState(canvas=canvas, recurrent_cls=new_cls),
+            local_patches=out_patches,
+        )

canvit_nnx-0.1.0/canvit_nnx/viewpoint.py

@@ -0,0 +1,78 @@
+"""Glimpse extraction via viewpoint sampling.
+
+Coordinate convention (shared with canvit_nnx.model):
+- Coordinates are in [-1, 1]²
+- (0, 0) is image center, (-1, -1) is top-left, (1, 1) is bottom-right
+- Axis order is (y, x) — matrix indexing convention, NOT Cartesian
+- scales: 1 = full scene, <1 = zoomed in
+
+Matches PyTorch F.grid_sample with align_corners=False.
+"""
+
+from __future__ import annotations
+
+import jax
+import jax.numpy as jnp
+
+from canvit_nnx.model import Viewpoint
+
+Array = jax.Array
+
+
+def sample_at_viewpoint(
+    spatial: Array,
+    viewpoint: Viewpoint,
+    glimpse_size_px: int,
+) -> Array:
+    """Bilinear crop from spatial tensor at viewpoint position.
+
+    Args:
+        spatial: [B, H, W, C] image or feature map (NHWC)
+        viewpoint: where to look (centers [B, 2], scales [B])
+        glimpse_size_px: output spatial size (square)
+
+    Returns:
+        [B, glimpse_size_px, glimpse_size_px, C] bilinearly sampled crop
+    """
+    B, H, W, C = spatial.shape
+    g = glimpse_size_px
+
+    # Sampling grid in [-1, 1] — cell centers
+    ys = (jnp.arange(g, dtype=jnp.float32) + 0.5) / g * 2 - 1
+    xs = (jnp.arange(g, dtype=jnp.float32) + 0.5) / g * 2 - 1
+    grid_y, grid_x = jnp.meshgrid(ys, xs, indexing="ij")
+    grid = jnp.stack([grid_y, grid_x], axis=-1)  # [g, g, 2]
+
+    # Scene-space positions: center + scale * grid
+    center = viewpoint.centers[:, None, None, :]  # [B, 1, 1, 2]
+    scale = viewpoint.scales[:, None, None, None]  # [B, 1, 1, 1]
+    scene_pos = center + scale * grid[None]  # [B, g, g, 2]
+
+    # [-1, 1] → pixel coords (align_corners=False: pixel = (grid + 1) * N / 2 - 0.5)
+    pixel_y = (scene_pos[..., 0] + 1) * H / 2 - 0.5
+    pixel_x = (scene_pos[..., 1] + 1) * W / 2 - 0.5
+
+    # Bilinear interpolation
+    y0 = jnp.floor(pixel_y).astype(jnp.int32)
+    x0 = jnp.floor(pixel_x).astype(jnp.int32)
+    y1 = y0 + 1
+    x1 = x0 + 1
+
+    y0c = jnp.clip(y0, 0, H - 1)
+    y1c = jnp.clip(y1, 0, H - 1)
+    x0c = jnp.clip(x0, 0, W - 1)
+    x1c = jnp.clip(x1, 0, W - 1)
+
+    wy = pixel_y - y0.astype(jnp.float32)
+    wx = pixel_x - x0.astype(jnp.float32)
+
+    b_idx = jnp.arange(B)[:, None, None]
+    val00 = spatial[b_idx, y0c, x0c]
+    val01 = spatial[b_idx, y0c, x1c]
+    val10 = spatial[b_idx, y1c, x0c]
+    val11 = spatial[b_idx, y1c, x1c]
+
+    return (val00 * (1 - wy[..., None]) * (1 - wx[..., None]) +
+            val01 * (1 - wy[..., None]) * wx[..., None] +
+            val10 * wy[..., None] * (1 - wx[..., None]) +
+            val11 * wy[..., None] * wx[..., None])

canvit_nnx-0.1.0/pyproject.toml

@@ -0,0 +1,47 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["canvit_nnx"]
+
+[tool.hatch.build.targets.sdist]
+exclude = ["throwaway/", "scripts/", "tests/"]
+
+[project]
+name = "canvit-nnx"
+version = "0.1.0"
+description = "CanViT (Canvas Vision Transformer) -- JAX / Flax NNX"
+readme = "README.md"
+license = "MIT"
+authors = [
+    { name = "Yohaï-Eliel Berreby", email = "me@yberreby.com" },
+]
+requires-python = ">=3.12"
+dependencies = [
+    "jax>=0.9",
+    "flax>=0.12",
+    "jaxlib>=0.9",
+    "safetensors",
+    "huggingface-hub",
+]
+
+[project.optional-dependencies]
+convert = [
+    "torch>=2.0",
+    "canvit-pytorch>=0.1.9",
+]
+verify = [
+    "canvit-nnx[convert]",
+    "pillow",
+]
+demo = [
+    "pillow",
+]
+dev = [
+    "canvit-nnx[convert,verify,demo]",
+    "torchvision",
+    "ipython",
+    "pytest",
+    "pypatree",
+]