This is an automated email from the ASF dual-hosted git repository.
cbalint13 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 4650887f6a [Relax][ONNX] Support align_corners in AffineGrid op
(#19864)
4650887f6a is described below
commit 4650887f6ab6db561246615fc6b17b5abfa81ed1
Author: Guan-Ming (Wesley) Chiu <[email protected]>
AuthorDate: Mon Jun 22 17:46:36 2026 +0800
[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.
---
include/tvm/relax/attrs/image.h | 13 ++++++++++
python/tvm/relax/frontend/onnx/onnx_frontend.py | 7 ++----
python/tvm/relax/op/image/image.py | 13 +++++-----
python/tvm/relax/transform/legalize_ops/image.py | 1 +
python/tvm/topi/image/grid_sample.py | 32 +++++++++++++++++-------
src/relax/op/image/resize.cc | 9 +++++--
src/relax/op/image/resize.h | 2 +-
tests/python/relax/test_frontend_onnx.py | 12 ++++-----
8 files changed, 58 insertions(+), 31 deletions(-)
diff --git a/include/tvm/relax/attrs/image.h b/include/tvm/relax/attrs/image.h
index eacbea7180..c9a7203740 100644
--- a/include/tvm/relax/attrs/image.h
+++ b/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
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py
b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index d550d3bc00..61f95e7130 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/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 @@ class AffineGrid(OnnxOpConverter):
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]))
diff --git a/python/tvm/relax/op/image/image.py
b/python/tvm/relax/op/image/image.py
index 323bfa74b5..29aa8457d2 100644
--- a/python/tvm/relax/op/image/image.py
+++ b/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))
diff --git a/python/tvm/relax/transform/legalize_ops/image.py
b/python/tvm/relax/transform/legalize_ops/image.py
index 42cb72f7c9..687e898a21 100644
--- a/python/tvm/relax/transform/legalize_ops/image.py
+++ b/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,
)
diff --git a/python/tvm/topi/image/grid_sample.py
b/python/tvm/topi/image/grid_sample.py
index 79032f41b3..cdfb7f4362 100644
--- a/python/tvm/topi/image/grid_sample.py
+++ b/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,6 +35,10 @@ 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
@@ -42,18 +46,28 @@ def affine_grid(data, target_shape):
"""
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)
diff --git a/src/relax/op/image/resize.cc b/src/relax/op/image/resize.cc
index beb89af087..b92167e031 100644
--- a/src/relax/op/image/resize.cc
+++ b/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);
diff --git a/src/relax/op/image/resize.h b/src/relax/op/image/resize.h
index 06a927d3a7..382a3a162b 100644
--- a/src/relax/op/image/resize.h
+++ b/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
diff --git a/tests/python/relax/test_frontend_onnx.py
b/tests/python/relax/test_frontend_onnx.py
index d3036a2547..82e3f997d2 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/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.
[email protected]("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],
