[Relax][ONNX] Support align_corners in AffineGrid op (#19864)

## Related Issue

closes #19690

## Why

ONNX AffineGrid carries an align_corners attribute, but the Relax op
ignored it and always produced the align_corners=1 grid.

## How

- Add align_corners field to AffineGridAttrs (mirrors GridSampleAttrs).
- Thread the flag through the op, legalize pass, and TOPI compute.
- Pass the ONNX attribute through in the frontend instead of dropping
it.

Author: guan404ming

include/tvm/relax/attrs/image.h

@@ -149,6 +149,19 @@ struct GridSampleAttrs : public AttrsNode {
   TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.GridSampleAttrs", GridSampleAttrs, AttrsNode);
 };  // struct GridSampleAttrs
 
+/*! \brief Attributes used in image affine_grid operator */
+struct AffineGridAttrs : public AttrsNode {
+  bool align_corners;
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<AffineGridAttrs>().def_ro(
+        "align_corners", &AffineGridAttrs::align_corners,
+        "If True, normalized grid coordinates map to corner pixels; otherwise to pixel centers.");
+  }
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AffineGridAttrs", AffineGridAttrs, AttrsNode);
+};  // struct AffineGridAttrs
+
 }  // namespace relax
 }  // namespace tvm
 

python/tvm/relax/frontend/onnx/onnx_frontend.py

@@ -3309,10 +3309,7 @@ class AffineGrid(OnnxOpConverter):
     def _impl_v20(cls, bb, inputs, attr, params):
         theta = inputs[0]  # [N, 2, 3] for 2D
         size = get_constant(inputs[1], params)  # [N, C, H, W] for 2D
-        align_corners = attr.get("align_corners", 0)
-
-        if align_corners != 1:
-            raise NotImplementedError("AffineGrid with align_corners=0 is not yet supported in TVM")
+        align_corners = bool(attr.get("align_corners", 0))
 
         # Extract size values
         if isinstance(size, relax.Constant):
@@ -3328,7 +3325,7 @@ def _impl_v20(cls, bb, inputs, attr, params):
         target_h, target_w = size_vals[2], size_vals[3]
 
         # Relax affine_grid outputs [N, 2, H, W]
-        grid = bb.emit(relax.op.image.affine_grid(theta, (target_h, target_w)))
+        grid = bb.emit(relax.op.image.affine_grid(theta, (target_h, target_w), align_corners))
         # Permute to ONNX convention [N, H, W, 2]
         return bb.emit(relax.op.permute_dims(grid, axes=[0, 2, 3, 1]))
 

python/tvm/relax/op/image/image.py

@@ -237,6 +237,7 @@ def grid_sample(
 def affine_grid(
     data: Expr,
     size: Expr | SizeLike,
+    align_corners: bool = True,
 ) -> Expr:
     """Generate a 2D sampling grid using an affine transformation matrix.
 
@@ -253,20 +254,18 @@ def affine_grid(
         The target output spatial shape (H, W). If a single integer or PrimExpr
         is provided, it is interpreted as a square output shape (size, size).
 
+    align_corners : bool
+        If True, normalized grid coordinates map to corner pixels; if False, to
+        pixel centers (the PyTorch / ONNX default).
+
     Returns
     -------
     result : relax.Expr
         The output grid tensor with shape [batch, 2, H, W].
-
-    Note
-    ----
-    Only `align_corners=True` is supported by this operator, matching the
-    behavior of the underlying TOPI implementation. When using this operator
-    via PyTorch or ONNX frontends, `align_corners=False` will be rejected.
     """
     if isinstance(size, int | PrimExpr):
         size = (size, size)
     if isinstance(size, tuple | list):
         size = ShapeExpr(size)
 
-    return cast(Expr, _ffi_api.affine_grid(data, size))
+    return cast(Expr, _ffi_api.affine_grid(data, size, align_corners))

python/tvm/relax/transform/legalize_ops/image.py

@@ -66,6 +66,7 @@ def _image_affine_grid(bb: BlockBuilder, call: Call) -> Expr:
         topi.image.affine_grid,
         call.args[0],
         target_shape=target_shape,
+        align_corners=call.attrs.align_corners,
     )
 
 

python/tvm/topi/image/grid_sample.py

@@ -20,7 +20,7 @@
 from tvm import te, tirx
 
 
-def affine_grid(data, target_shape):
+def affine_grid(data, target_shape, align_corners=True):
     """affine_grid operator that generates 2D sampling grid.
 
     This operation is described in https://arxiv.org/pdf/1506.02025.pdf. It generates a uniform
@@ -35,25 +35,39 @@ def affine_grid(data, target_shape):
     target_shape: list/tuple of two int
         Specifies the output shape (H, W).
 
+    align_corners : bool
+        If True, normalized coordinates map to corner pixels; if False, to pixel centers
+        (the PyTorch / ONNX default).
+
     Returns
     -------
     Output : tvm.Tensor
         4-D with shape [batch, 2, target_height, target_width]
     """
     assert target_shape is not None
     assert len(target_shape) == 2
-    assert target_shape[0] > 1 and target_shape[1] > 1, (
-        "target height/width should be greater than 1"
-    )
+    if align_corners:
+        assert target_shape[0] > 1 and target_shape[1] > 1, (
+            "target height/width should be greater than 1 when align_corners is True"
+        )
 
     dtype = data.dtype
-    y_step = tirx.const((2.0 - 1e-7) / (target_shape[0] - 1), dtype=dtype)
-    x_step = tirx.const((2.0 - 1e-7) / (target_shape[1] - 1), dtype=dtype)
-    start = tirx.const(-1.0, dtype=dtype)
+    height, width = target_shape[0], target_shape[1]
+    if align_corners:
+        y_step = tirx.const((2.0 - 1e-7) / (height - 1), dtype=dtype)
+        x_step = tirx.const((2.0 - 1e-7) / (width - 1), dtype=dtype)
+        y_start = tirx.const(-1.0, dtype=dtype)
+        x_start = tirx.const(-1.0, dtype=dtype)
+    else:
+        # Pixel centers: coordinate i maps to (2 * i + 1) / size - 1.
+        y_step = tirx.const(2.0 / height, dtype=dtype)
+        x_step = tirx.const(2.0 / width, dtype=dtype)
+        y_start = tirx.const(-1.0 + 1.0 / height, dtype=dtype)
+        x_start = tirx.const(-1.0 + 1.0 / width, dtype=dtype)
 
     def _compute(n, dim, i, j):
-        y = start + i * y_step
-        x = start + j * x_step
+        y = y_start + i * y_step
+        x = x_start + j * x_step
         return data[n, dim, 0] * x + data[n, dim, 1] * y + data[n, dim, 2]
 
     oshape = (data.shape[0], len(target_shape), *target_shape)

src/relax/op/image/resize.cc

@@ -354,11 +354,15 @@ TVM_REGISTER_OP("relax.image.grid_sample")
 
 /* relax.image.affine_grid */
 
-Expr affine_grid(Expr data, Expr size) {
+Expr affine_grid(Expr data, Expr size, bool align_corners) {
+  ffi::ObjectPtr<AffineGridAttrs> attrs = ffi::make_object<AffineGridAttrs>();
+  attrs->align_corners = align_corners;
   static const Op& op = Op::Get("relax.image.affine_grid");
-  return Call(op, {std::move(data), std::move(size)}, Attrs(), {});
+  return Call(op, {std::move(data), std::move(size)}, Attrs(attrs), {});
 }
 
+TVM_FFI_STATIC_INIT_BLOCK() { AffineGridAttrs::RegisterReflection(); }
+
 TVM_FFI_STATIC_INIT_BLOCK() {
   namespace refl = tvm::ffi::reflection;
   refl::GlobalDef().def("relax.op.image.affine_grid", affine_grid);
@@ -438,6 +442,7 @@ TVM_REGISTER_OP("relax.image.affine_grid")
     .set_num_inputs(2)
     .add_argument("data", "Tensor", "The input affine matrix tensor.")
     .add_argument("size", "Shape", "The target output shape (H, W).")
+    .set_attrs_type<AffineGridAttrs>()
     .set_attr<FInferType>("FInferType", InferTypeAffineGrid)
     .set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kFollow)
     .set_attr<bool>("FPurity", true);

src/relax/op/image/resize.h

@@ -49,7 +49,7 @@ Expr grid_sample(Expr data, Expr grid, ffi::String method, ffi::String layout,
                  ffi::String padding_mode, bool align_corners);
 
 /*! \brief Image affine_grid operator. */
-Expr affine_grid(Expr data, Expr size);
+Expr affine_grid(Expr data, Expr size, bool align_corners);
 
 }  // namespace relax
 }  // namespace tvm

tests/python/relax/test_frontend_onnx.py

@@ -5555,13 +5555,11 @@ def test_nms_score_threshold():
         )
 
 
-def test_affine_grid():
-    affine_grid_node = helper.make_node(
-        "AffineGrid",
-        ["theta", "size"],
-        ["grid"],
-        align_corners=1,
-    )
+# align_corners=None omits the attribute, exercising the ONNX default of 0.
+@pytest.mark.parametrize("align_corners", [None, 0, 1])
+def test_affine_grid(align_corners):
+    attrs = {} if align_corners is None else {"align_corners": align_corners}
+    affine_grid_node = helper.make_node("AffineGrid", ["theta", "size"], ["grid"], **attrs)
 
     graph = helper.make_graph(
         [affine_grid_node],