Utonia — JAX / Flax NNX weights
[ GitHub Repo ]
JAX conversion of the Utonia weights. Utonia is from the paper Utonia: Toward One Encoder for All Point Clouds (arXiv:2603.03283). It is a Point Transformer V3 encoder, pretrained with self-supervision on indoor RGB-D, outdoor LiDAR, remote sensing, object CAD and video-lifted point clouds.
The weights are meant for use with utonia-jax, a Flax NNX port of the Utonia inference code. It does not need spconv, torch_scatter or flash-attn.
Files
| Path | Contents |
|---|---|
utonia/model.safetensors | Encoder weights (137.25M params, float32), converted from utonia.pth |
utonia/config.json | Model config from the original checkpoint |
utonia_linear_prob_head_sc/model.safetensors | ScanNet-20 linear-probe segmentation head (Linear(1386 → 20)), converted from utonia_linear_prob_head_sc.pth |
utonia_linear_prob_head_sc/config.json | Seg-head config |
Changes from the original
Converted from the PyTorch checkpoints in Pointcept/Utonia. There was no retraining or fine-tuning, and every value is unchanged. Only the format and layout changed:
.pth(pickled state dict) was converted to.safetensors, float32.- Tensor names are Flax NNX state paths (e.g.
enc/3/blocks/0/attn/qkv/kernel). - Linear weights are transposed to
(in, out). - spconv
SubMConv3dkernels are rearranged from(C_out, 3, 3, 3, C_in)to(27, C_in, C_out).
Numerical parity with the reference PyTorch implementation, on the ScanNet demo scene:
| Attention precision | Feature rel. error | Seg argmax agreement |
|---|---|---|
| fp32 | ≤ 3.4e-4 | 99.99% |
| bf16 (default, like upstream flash-attn) | ≤ 1.3e-3 | 99.98% |
Usage
pip install git+https://github.com/abhaybd/Utonia-Jax
import jax
import utonia_jax as utonia
model = utonia.load("utonia") # downloads from this repo on first use
head = utonia.load_seg_head() # ScanNet-20 linear probe
point = utonia.data.load("sample1")
point = utonia.transform.default(scale=0.5)(point)
batch = utonia.prepare(point, model.config)
@jax.jit
def encode(batch):
return utonia.upcast(model(batch), batch) # (M, 1386) per-voxel features
feat = encode(batch)[: int(batch.stages[0].num_points)]
pred = head(feat).argmax(-1)[point["inverse"]] # per-original-point labels
License
CC-BY-NC-4.0, inherited from the original Pointcept/Utonia weights. The weights may be used for non-commercial purposes only. The utonia-jax code is licensed separately under Apache 2.0.
All credit for the model goes to the Utonia authors.
Citation
@inproceedings{zhang2026utonia,
title={Utonia: Toward One Encoder for All Point Clouds},
author={Yujia Zhang and Xiaoyang Wu and Yunhan Yang and Xianzhe Fan and Han Li and Yuechen Zhang and Zehao Huang and Naiyan Wang and Hengshuang Zhao},
booktitle={ICML},
year={2026},
}
@misc{pointcept2023,
title={Pointcept: A Codebase for Point Cloud Perception Research},
author={Pointcept Contributors},
howpublished={\url{https://github.com/Pointcept/Pointcept}},
year={2023}
}