abhaybd/Utonia-Jax

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 SubMConv3d kernels 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}
}