[Bug][Relax][ONNX] AffineGrid with `align_corners=0` (the ONNX default!) is unsupported

[wuyii8941]

Expected behavior

ONNX AffineGrid defaults to align_corners=0. Both ONNX Runtime and PyTorch implement the align_corners=0 ("center of pixel") sampling convention.

Actual behavior

The TVM converter rejects align_corners=0:

NotImplementedError: AffineGrid with align_corners=0 is not yet supported in TVM

Because the ONNX default is 0 (the attribute is normally omitted from the graph), every ONNX model that uses AffineGrid without explicitly setting align_corners=1 fails to import.

Reproduction

import numpy as np
import onnx
from onnx import helper, TensorProto, numpy_helper
import onnxruntime as ort
from tvm.relax.frontend.onnx import from_onnx

theta = helper.make_tensor_value_info("theta", TensorProto.FLOAT, [2, 2, 3])
Y = onnx.ValueInfoProto(); Y.name = "Y"
size = numpy_helper.from_array(np.array([2, 1, 4, 5], dtype=np.int64), "size")
# align_corners defaults to 0 — attribute omitted from node, as ONNX exporters do
node = helper.make_node("AffineGrid", ["theta", "size"], ["Y"])
g = helper.make_graph([node], "g", [theta], [Y], initializer=[size])
m = helper.make_model(g, opset_imports=[helper.make_opsetid("", 20)])

theta_v = np.random.randn(2, 2, 3).astype(np.float32)
print("ORT shape:", ort.InferenceSession(m.SerializeToString()).run(None, {"theta": theta_v})[0].shape)
# ORT shape: (2, 4, 5, 2)

inf = onnx.shape_inference.infer_shapes(m)
mod = from_onnx(inf)  # NotImplementedError: AffineGrid with align_corners=0 ...

Root cause

python/tvm/relax/frontend/onnx/onnx_frontend.py, AffineGrid._impl_v20:

align_corners = attr.get("align_corners", 0)

if align_corners != 1:
    raise NotImplementedError("AffineGrid with align_corners=0 is not yet supported in TVM")

The underlying relax.op.image.affine_grid (and corresponding TOPI implementation) currently only implements the align_corners=1 formula.

For 2D, the difference is the source-coordinate mapping:

• align_corners=1: x = (-1, +1) ⟼ (0, W-1) — corners aligned
• align_corners=0: x = (-1, +1) ⟼ (-0.5, W-0.5) — pixel centers aligned

This is the standard PyTorch convention and is what ONNX Runtime implements.

Suggested fix

Plumb align_corners into relax.op.image.affine_grid / topi.image.affine_grid so both modes are supported, mirroring how image.resize2d already takes coordinate_transformation_mode.

Minimal short-term workaround would be to emit the coordinate offsets in the frontend itself (compute the grid via composing arange + scale + translate with align_corners=0 math), but the cleaner fix is in the op.

Impact

Every ONNX export that uses AffineGrid without explicitly setting align_corners=1 (i.e., almost all of them) fails to import. This blocks Spatial Transformer Networks and any geometric-warping model exported from PyTorch with default arguments.

Environment

• TVM: latest main (commit b172d5e)
• Python: 3.11
• ONNX Runtime: 1.24.4

cc @KJlaccHoeUM9l @junrushao