"""DINOv3 ViT-L/16 features on trn2 with vit-encoder-neuron-kernels.

    NEURON_LOGICAL_NC_CONFIG=2 python dinov3_features.py
"""
import types
import torch
from kernels import get_kernel
from transformers import AutoModel

k = get_kernel("aws-neuron/vit-encoder-neuron-kernels", version=1, trust_remote_code=True)
model = AutoModel.from_pretrained("facebook/dinov3-vitl16-pretrain-lvd1689m", dtype=torch.float32).eval()
ref_model = AutoModel.from_pretrained("facebook/dinov3-vitl16-pretrain-lvd1689m", dtype=torch.float32).eval()
model = model.to("neuron")
model.forward = types.MethodType(k.NeuronDINOv3ViTModel.forward, model)
model = torch.compile(model, backend="neuron")

torch.manual_seed(0)
pixel_values = torch.randn(4, 3, 224, 224)
with torch.no_grad():
    out = model(pixel_values.to("neuron")).last_hidden_state.cpu()
    ref = ref_model(pixel_values).last_hidden_state
print("features", tuple(out.shape), "rel_l2 vs fp32 CPU:", ((out - ref).norm() / ref.norm()).item())