[Bug][Relax][ONNX] AffineGrid 3D output (`size=[N,C,D,H,W]`) is rejected with `ValueError`

[wuyii8941]

Expected behavior

ONNX AffineGrid (opset 20) supports both 2D output (size=[N,C,H,W], theta shape [N, 2, 3]) and 3D output (size=[N,C,D,H,W], theta shape [N, 3, 4]). PyTorch's F.affine_grid and ONNX Runtime both handle 3D.

Actual behavior

The TVM converter raises:

ValueError: Only 2D AffineGrid (size=[N,C,H,W]) is supported

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, 3, 4])
Y = onnx.ValueInfoProto(); Y.name = "Y"
size = numpy_helper.from_array(np.array([2, 1, 3, 4, 5], dtype=np.int64), "size")
node = helper.make_node("AffineGrid", ["theta", "size"], ["Y"], align_corners=1)
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, 3, 4).astype(np.float32)
print("ORT shape:", ort.InferenceSession(m.SerializeToString()).run(None, {"theta": theta_v})[0].shape)
# ORT shape: (2, 3, 4, 5, 3)

inf = onnx.shape_inference.infer_shapes(m)
mod = from_onnx(inf)  # ValueError: Only 2D AffineGrid is supported

Root cause

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

# Only 2D is supported: size = [N, C, H, W]
if len(size_vals) != 4:
    raise ValueError("Only 2D AffineGrid (size=[N,C,H,W]) is supported")
target_h, target_w = size_vals[2], size_vals[3]

grid = bb.emit(relax.op.image.affine_grid(theta, (target_h, target_w)))
return bb.emit(relax.op.permute_dims(grid, axes=[0, 2, 3, 1]))

The underlying relax.op.image.affine_grid and topi.image.affine_grid only target 2D. To support 3D, either:

• Extend affine_grid TOPI / Relax op to accept a 3D target shape (preferable, matches PyTorch behavior), or
• At minimum, emit a NotImplementedError explicitly mentioning 3D as the limitation (current ValueError reads as if 3D is invalid input rather than missing implementation).

Suggested fix

Add a 3D code path when len(size_vals) == 5. PyTorch's F.affine_grid for 3D produces a [N, D, H, W, 3] output. Implementation pattern:

elif len(size_vals) == 5:
    target_d, target_h, target_w = size_vals[2], size_vals[3], size_vals[4]
    grid = bb.emit(relax.op.image.affine_grid_3d(theta, (target_d, target_h, target_w)))
    return bb.emit(relax.op.permute_dims(grid, axes=[0, 2, 3, 4, 1]))

(requires adding a 3D affine_grid variant, or generalizing the existing op).

Impact

3D Spatial Transformer Networks (STNs), volumetric warping models, and any opset-20+ ONNX export that uses 3D affine grids cannot be imported.

Environment

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

cc @KJlaccHoeUM9l @junrushao