[Tracking Issue] Add missing vision/image operators

[tlopex]

Background

Since we phased out Relay in early 2025, several vision/image operators that Relay supported are missing from Relax, blocking deployment of common detection and segmentation models (Faster R-CNN, Mask R-CNN, Spatial Transformer Networks, etc.).

Currently Relax only has 3 vision/image ops: all_class_non_max_suppression, resize2d, grid_sample.

Scope

Each operator needs:

1. TOPI compute (python/tvm/topi/) — if not already present
2. Relax op registration — C++ attrs + struct info inference (src/relax/op/, include/tvm/relax/attrs/) + Python wrapper (python/tvm/relax/op/)
3. Legalization (python/tvm/relax/transform/legalize_ops/)
4. Frontend integration — update ONNX/PyTorch/TFLite frontends to emit the new op
5. Tests — op-level unit tests + frontend integration tests

Steps 1–3 (Relax op) and step 4 (frontend) can be done in separate PRs or combined — contributors can choose based on scope.

Operators

Tier 1 — Unblocks mainstream detection/segmentation models

• roi_align — Core op for Faster R-CNN, Mask R-CNN, Cascade R-CNN. TOPI was removed during te.schedule phase-out, needs full reimplementation.
• affine_grid — Generates sampling grid for Spatial Transformer Networks. TOPI already exists (topi.image.affine_grid), only needs Relax op + legalization. Pairs with existing grid_sample.

Tier 2 — Fills domain gaps and fixes broken paths

• resize3d — 3D volume resize for medical imaging (CT/MRI) and video. TOPI already exists (topi.image.resize3d), only needs Relax op + legalization. Note: ONNX frontend currently works around this by calling TOPI directly.
• get_valid_counts — Score-based bounding box filtering. Filters out low-score boxes and returns valid count per batch. Current TOPI is a no-op stub; needs full implementation.
• non_max_suppression (classic) — Flexible single-class NMS. Performs IoU-based suppression on filtered boxes from get_valid_counts. TOPI was removed; needs reimplementation. Complements existing all_class_non_max_suppression for custom post-processing pipelines.
• multibox_transform_loc — Decodes bounding box predictions using anchor priors + predicted offsets, with score thresholding. Only needed by TFLite DETECTION_POSTPROCESS. No existing TOPI or Relax op; needs full implementation.

Tier 3 — Frontend integration

• ONNX RoiAlign: add relax.op.vision.roi_align op definition + legalization, implement converter class, uncomment registration, add tests
• ONNX AffineGrid: add relax.op.image.affine_grid op definition + legalization, implement converter class, add tests
• ONNX Resize 5-D: add relax.op.image.resize3d op definition + legalization, replace bb.emit_te workaround in converter with relax.op.image.resize3d, add tests
• ONNX GridSample: implement converter class (relax op already exists), uncomment registration, add tests
• PyTorch torchvision.ops.roi_align: register converter (depends on relax.op.vision.roi_align), add tests
• PyTorch torch.nn.functional.affine_grid: register converter (depends on relax.op.image.affine_grid), add tests
• PyTorch torch.nn.functional.interpolate 3D mode: add 5-D branch in converter using relax.op.image.resize3d (depends on resize3d op), add tests
• TFLite NON_MAX_SUPPRESSION_V5: verify/add get_valid_counts + non_max_suppression relax ops, implement converter, add tests
• TFLite DETECTION_POSTPROCESS: implement converter (anchor decoding + NMS), depends on same relax ops as above, add tests

Implementation reference

• Existing vision op pattern: src/relax/op/vision/nms.cc + python/tvm/relax/op/vision/nms.py
• Existing image op pattern: src/relax/op/image/resize.cc + python/tvm/relax/op/image/image.py
• Legalization: python/tvm/relax/transform/legalize_ops/vision.py, image.py
• roi_align reference implementation for testing: python/tvm/topi/testing/roi_align_python.py
• Historical Relay implementations as design reference (Relax has a different IR design, so these should not be copied directly): v0.19.0 tag

affine_grid and resize3d are good first issues — TOPI already exists, just follow the existing resize2d / grid_sample pattern. If you're interested in contributing, please comment below to claim a task before starting work.

[Dayuxiaoshui]

Hi @tlopex, I'd like to take on resize3d.

[Aharrypotter]

Hi @tlopex, I’d like to claim affine_grid, including the Relax op + legalization and the PyTorch frontend integration for torch.nn.functional.affine_grid.

[LudovicoYIN]

Hi @tlopex, I’d like to take on roi_align.

[nirdesh17]

Hi @tlopex , I’d like to work on ONNX frontend integration for GridSample.

[OmarAzizi]

Hi @tlopex, I'd like to claim the ONNX frontend integration for Resize 5-D → resize3d. I'll start working on it now and open a draft PR in the meantime while @Dayuxiaoshui finishes the resize3d Relax op, so the two can land together.

[OmarAzizi]

@tlopex @nirdesh17 I apologize for the overlap. I saw that nirdesh17 had
claimed GridSample, but I was excited to contribute and couldn't wait to dive in
as a beginner. If nirdesh17 has already started working on it, I'm happy to close
my PR and defer to them. Sorry about that!

[nirdesh17]

@tlopex @nirdesh17 I apologize for the overlap. I saw that nirdesh17 had claimed GridSample, but I was excited to contribute and couldn't wait to dive in as a beginner. If nirdesh17 has already started working on it, I'm happy to close my PR and defer to them. Sorry about that!

@OmarAzizi No worries at all! I had just started exploring it and was still understanding things, so please go ahead with your PR. I’ll follow along and learn from it

[nirdesh17]

Hi @tlopex, @Aharrypotter — I see affine_grid Relax op and PyTorch frontend are being worked on. I’d like to take up the ONNX frontend integration for affine_grid once the backend is ready. Please let me know if that’s okay.

[Aharrypotter]

Hi @nirdesh17, thanks for your interest! The ONNX frontend integration for AffineGrid is actually already included in my PR #18933 — it covers the full Relax op + legalization + PyTorch frontend + ONNX frontend (opset 20) as a single package.

If you're looking for other tasks to pick up, there are still several unclaimed items on this issue:

• PyTorch interpolate 3D mode (Tier 3) — depends on resize3d ([Relax][ONNX] Add image.resize3d op and wire 5D Resize #18931), frontend converter needed

Happy to help if you have questions about any of these!

[Aharrypotter]

Hi @tlopex, I'd like to claim get_valid_counts + non_max_suppression (classic) (Tier 2). I'll start with the TOPI implementation for both, then build up the Relax op registration, legalization, and tests.