From f61e9b82d15c5e19afabbdc82835de335f13742e Mon Sep 17 00:00:00 2001 From: Asthenia0412 Date: Fri, 16 Oct 2026 05:02:02 +0000 Subject: [PATCH 1/4] [Relax][Op] Add missing TOPI-backed operators: correlation, lrn, dilate, dilation2d Add 4 new Relax operators that were previously only available in TOPI: - relax.nn.correlation: Cross-correlation between two feature maps (NCHW layout) - relax.nn.lrn: Local Response Normalization (AlexNet-style across channels) - relax.nn.dilate: Upsample with holes (insert zeros between elements) - relax.image.dilation2d: Morphological dilation (NCHW and NHWC layouts) Each operator includes: - C++ Attrs class with TVM_FFI_DECLARE_OBJECT_INFO_FINAL macro - Python Relax op definition with proper attrs and docstrings - TOPI legalization in legalize_ops/ - Export in __init__.py modules These operators fill gaps between TOPI's rich operator set and Relax's op surface, enabling easier frontend integration (ONNX/TFLite/Torch) for models that use these common but previously unsupported ops. --- include/tvm/relax/op/image.h | 12 + include/tvm/relax/op/nn.h | 35 +++ python/tvm/relax/op/image/__init__.py | 2 + python/tvm/relax/op/image/image.py | 98 ++++++++ python/tvm/relax/op/nn/__init__.py | 6 + python/tvm/relax/op/nn/nn.py | 212 ++++++++++++++++++ .../tvm/relax/transform/legalize_ops/image.py | 23 ++ python/tvm/relax/transform/legalize_ops/nn.py | 50 +++++ src/relax/op/image/resize.cc | 16 ++ src/relax/op/nn/nn.cc | 38 ++++ 10 files changed, 492 insertions(+) diff --git a/include/tvm/relax/op/image.h b/include/tvm/relax/op/image.h index 9aa1beb98bf2..c649afd4ba6b 100644 --- a/include/tvm/relax/op/image.h +++ b/include/tvm/relax/op/image.h @@ -80,6 +80,18 @@ struct AffineGridAttrs : public AttrsNode { TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AffineGridAttrs", AffineGridAttrs, AttrsNode); }; // struct AffineGridAttrs +/*! \brief Attributes used in image dilation2d operator */ +struct Dilation2DAttrs : public AttrsNode { + ffi::Array strides; + ffi::Array padding; + ffi::Array dilations; + ffi::String data_layout; + ffi::Optional out_dtype; + + static void RegisterReflection(); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Dilation2DAttrs", Dilation2DAttrs, AttrsNode); +}; // struct Dilation2DAttrs + } // namespace relax } // namespace tvm diff --git a/include/tvm/relax/op/nn.h b/include/tvm/relax/op/nn.h index 3a147ac54a6a..997a6620773f 100644 --- a/include/tvm/relax/op/nn.h +++ b/include/tvm/relax/op/nn.h @@ -339,6 +339,41 @@ struct PixelShuffleAttrs : public AttrsNode { TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.PixelShuffleAttrs", PixelShuffleAttrs, AttrsNode); }; +/*! \brief Attributes used for the correlation operator */ +struct CorrelationAttrs : public AttrsNode { + int kernel_size; + int max_displacement; + int stride1; + int stride2; + int padding; + bool is_multiply; + ffi::String data_layout; + + static void RegisterReflection(); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.CorrelationAttrs", CorrelationAttrs, AttrsNode); +}; // struct CorrelationAttrs + +/*! \brief Attributes used for the local response normalization (LRN) operator */ +struct LRNAattrs : public AttrsNode { + int size; + int axis; + double alpha; + double beta; + double bias; + + static void RegisterReflection(); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.LRNAattrs", LRNAattrs, AttrsNode); +}; // struct LRNAattrs + +/*! \brief Attributes used for the dilate (upsample with holes) operator */ +struct DilateAttrs : public AttrsNode { + ffi::Array strides; + double dilation_value; + + static void RegisterReflection(); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.DilateAttrs", DilateAttrs, AttrsNode); +}; // struct DilateAttrs + } // namespace relax } // namespace tvm diff --git a/python/tvm/relax/op/image/__init__.py b/python/tvm/relax/op/image/__init__.py index daa645b5e607..32be1587db8d 100644 --- a/python/tvm/relax/op/image/__init__.py +++ b/python/tvm/relax/op/image/__init__.py @@ -19,7 +19,9 @@ from .image import ( Resize2DAttrs, + Dilation2DAttrs, affine_grid, + dilation2d, grid_sample, resize2d, resize3d, diff --git a/python/tvm/relax/op/image/image.py b/python/tvm/relax/op/image/image.py index f2a159497600..abe06ade4010 100644 --- a/python/tvm/relax/op/image/image.py +++ b/python/tvm/relax/op/image/image.py @@ -317,3 +317,101 @@ def affine_grid( ty=ty, ), ) + + +@tvm_ffi.register_object("relax.attrs.Dilation2DAttrs") +class Dilation2DAttrs(Attrs): + """Attributes for image.dilation2d""" + + +def dilation2d( + input: Expr, + filter: Expr, + strides: int | tuple[int, int] = (1, 1), + padding: int | tuple[int, int, int, int] | str = 0, + dilations: int | tuple[int, int] = (1, 1), + data_layout: str = "NCHW", + out_dtype: str | DataType | None = None, + *, + ty=None, + loc: Location = UNKNOWN_LOC, +) -> Expr: + r"""Morphological Dilation 2D operator. + + This operator computes the morphological dilation of the input image + with the given filter (structuring element). For each position, the + output is the maximum of `input + filter` over the filter window. + + .. math:: + + output[n, c, h, w] = \max_{f_h, f_w} + \left( input[n, c, h' + d_h \cdot f_h, w' + d_w \cdot f_w] + + filter[c, f_h, f_w] \right) + + where :math:`h', w'` are the top-left positions of the filter window + and :math:`d_h, d_w` are the dilation factors. + + Parameters + ---------- + input : relax.Expr + 4-D input tensor. With NCHW layout: [batch, in_channel, in_height, in_width]. + With NHWC layout: [batch, in_height, in_width, in_channel]. + + filter : relax.Expr + 3-D filter tensor with shape [in_channel, filter_height, filter_width]. + + strides : Union[int, Tuple[int, int]] + Stride size for the filter window, or [stride_h, stride_w]. + + padding : Union[int, Tuple[int, ...]] + Padding size. Can be an int for symmetric padding, + [pad_top, pad_left, pad_bottom, pad_right], or 2-tuple + [pad_h, pad_w] which is expanded to 4-tuple. + + dilations : Union[int, Tuple[int, int]] + Dilation size for the filter, or [dilation_h, dilation_w]. + + data_layout : str + Layout of the input data. Supported: "NCHW", "NHWC". + + out_dtype : Optional[str | DataType] + Specifies the output data type. If None, uses the input dtype. + + Returns + ------- + result : relax.Expr + The dilated output tensor. + """ + if isinstance(strides, int): + strides = (strides, strides) + if isinstance(dilations, int): + dilations = (dilations, dilations) + + if isinstance(padding, int): + padding = (padding, padding, padding, padding) + elif isinstance(padding, (tuple, list)): + if len(padding) == 2: + padding = (padding[0], padding[1], padding[0], padding[1]) + elif len(padding) == 4: + padding = tuple(padding) + else: + raise ValueError( + f"padding must be int, 2-tuple, or 4-tuple, got {len(padding)}-tuple" + ) + else: + raise TypeError(f"padding must be int, tuple, or list, got {type(padding)}") + + return _Call( + "relax.image.dilation2d", + [input, filter], + attrs=_make_attrs( + "relax.attrs.Dilation2DAttrs", + strides=strides, + padding=padding, + dilations=dilations, + data_layout=data_layout, + out_dtype=(out_dtype.dtype if isinstance(out_dtype, _PrimType) else out_dtype), + ), + ty=ty, + loc=loc, + ) # type: ignore diff --git a/python/tvm/relax/op/nn/__init__.py b/python/tvm/relax/op/nn/__init__.py index 9304b1dcbbf6..6af70b88521e 100644 --- a/python/tvm/relax/op/nn/__init__.py +++ b/python/tvm/relax/op/nn/__init__.py @@ -44,6 +44,9 @@ GroupNormAttrs, RMSNormAttrs, NLLLossAttrs, + CorrelationAttrs, + LRNAattrs, + DilateAttrs, adaptive_avg_pool1d, adaptive_avg_pool2d, adaptive_avg_pool3d, @@ -61,7 +64,9 @@ conv2d_transpose, conv3d, conv3d_transpose, + correlation, cross_entropy_with_logits, + dilate, dropout, gelu, gelu_tanh, @@ -70,6 +75,7 @@ layer_norm, leakyrelu, log_softmax, + lrn, max_pool1d, max_pool2d, max_pool3d, diff --git a/python/tvm/relax/op/nn/nn.py b/python/tvm/relax/op/nn/nn.py index 2ea810d1908c..ed8653882f6d 100644 --- a/python/tvm/relax/op/nn/nn.py +++ b/python/tvm/relax/op/nn/nn.py @@ -3007,3 +3007,215 @@ def batch_flatten(data: Expr, *, ty=None, loc: Location = UNKNOWN_LOC) -> Expr: The flattened result with shape `(batch_size, flattened_features)`. """ return _Call("relax.nn.batch_flatten", [data], loc=loc, ty=ty) # type: ignore + + +@tvm_ffi.register_object("relax.attrs.CorrelationAttrs") +class CorrelationAttrs(Attrs): + """Attributes for nn.correlation""" + + +def correlation( + data1: Expr, + data2: Expr, + kernel_size: int = 1, + max_displacement: int = 1, + stride1: int = 1, + stride2: int = 1, + padding: int = 0, + is_multiply: bool = True, + data_layout: str = "NCHW", + *, + ty=None, + loc: Location = UNKNOWN_LOC, +) -> Expr: + r"""Correlation operator. + + This operator computes correlation between two feature maps. + It is commonly used in optical flow networks. + + .. math:: + + out[b, q, i, j] = \frac{1}{kernel\_size^2 \times C} + \sum_{c=0}^{C-1} \sum_{ry=0}^{kernel\_size-1} \sum_{rx=0}^{kernel\_size-1} + f(data1[b, c, y1+ry, x1+rx], data2[b, c, y2+ry, x2+rx]) + + where :math:`f(x, y) = x \times y` if `is_multiply` is True, + otherwise :math:`f(x, y) = |x - y|`. + + Parameters + ---------- + data1 : relax.Expr + 4-D input tensor with shape [batch, channel, height, width]. + + data2 : relax.Expr + 4-D input tensor with the same shape as data1. + + kernel_size : int + Kernel size for correlation, must be a positive odd number. + + max_displacement : int + Max displacement of correlation. + + stride1 : int + Stride for sampling data1 positions. + + stride2 : int + Stride for sampling data2 within the neighborhood around data1. + + padding : int + Padding size for both height and width dimensions. + + is_multiply : bool + If True, compute element-wise product; otherwise compute absolute difference. + + data_layout : str + Layout of the input data. Currently only "NCHW" is supported. + + Returns + ------- + result : relax.Expr + The correlation result with shape [batch, out_channel, out_height, out_width]. + """ + if data_layout != "NCHW": + raise NotImplementedError( + f"correlation currently only supports NCHW layout, got {data_layout}" + ) + return _Call( + "relax.nn.correlation", + [data1, data2], + attrs=_make_attrs( + "relax.attrs.CorrelationAttrs", + kernel_size=kernel_size, + max_displacement=max_displacement, + stride1=stride1, + stride2=stride2, + padding=padding, + is_multiply=is_multiply, + data_layout=data_layout, + ), + loc=loc, + ty=ty, + ) # type: ignore + + +@tvm_ffi.register_object("relax.attrs.LRNAattrs") +class LRNAattrs(Attrs): + """Attributes for nn.lrn (local response normalization)""" + + +def lrn( + data: Expr, + size: int = 5, + axis: int = 1, + alpha: float = 0.0001, + beta: float = 0.75, + bias: float = 2.0, + *, + ty=None, + loc: Location = UNKNOWN_LOC, +) -> Expr: + r"""Local Response Normalization. + + This operator performs across-channel local response normalization + on the input data. It follows the formula from AlexNet: + + .. math:: + + output^i_{x,y} = \frac{data^i_{x,y}} + {\left(bias + \frac{\alpha}{size} + \sum_{j=\max(0, i-\lfloor size/2 \rfloor)} + ^{\min(N-1, i+\lfloor size/2 \rfloor)} + (data^j_{x,y})^2\right)^{\beta}} + + Parameters + ---------- + data : relax.Expr + 4-D input tensor with shape [batch, channel, height, width]. + + size : int + Normalization window size. + + axis : int + The axis along which normalization is applied. Default is 1 (channel axis). + + alpha : float + The scaling factor. + + beta : float + The exponent. + + bias : float + Offset to avoid dividing by zero. + + Returns + ------- + result : relax.Expr + The normalized output tensor with the same shape as input. + """ + return _Call( + "relax.nn.lrn", + [data], + attrs=_make_attrs( + "relax.attrs.LRNAattrs", + size=size, + axis=axis, + alpha=alpha, + beta=beta, + bias=bias, + ), + loc=loc, + ty=ty, + ) # type: ignore + + +@tvm_ffi.register_object("relax.attrs.DilateAttrs") +class DilateAttrs(Attrs): + """Attributes for nn.dilate""" + + +def dilate( + data: Expr, + strides: list[int], + dilation_value: float = 0.0, + *, + ty=None, + loc: Location = UNKNOWN_LOC, +) -> Expr: + """Dilate (upsample with holes) the input tensor. + + This operation inserts `strides[i] - 1` zeros between each element + along dimension `i`. The output shape is: + + output_shape[i] = (input_shape[i] - 1) * strides[i] + 1 + + Parameters + ---------- + data : relax.Expr + Input tensor of any shape. + + strides : list[int] + Dilation stride for each dimension. `1` means no dilation. + Must have the same length as data's number of dimensions. + + dilation_value : float + The value used to fill the inserted positions. Default is 0.0. + + Returns + ------- + result : relax.Expr + The dilated output tensor. + """ + if not isinstance(strides, (list, tuple)): + raise TypeError(f"strides must be a list or tuple of ints, got {type(strides)}") + strides = [int(s) for s in strides] + return _Call( + "relax.nn.dilate", + [data], + attrs=_make_attrs( + "relax.attrs.DilateAttrs", + strides=strides, + dilation_value=dilation_value, + ), + loc=loc, + ty=ty, + ) # type: ignore diff --git a/python/tvm/relax/transform/legalize_ops/image.py b/python/tvm/relax/transform/legalize_ops/image.py index 1b98fd80ca28..9602a6948fee 100644 --- a/python/tvm/relax/transform/legalize_ops/image.py +++ b/python/tvm/relax/transform/legalize_ops/image.py @@ -87,3 +87,26 @@ def _image_resize3d(bb: BlockBuilder, call: Call) -> Expr: bicubic_exclude=call.attrs.cubic_exclude, extrapolation_value=call.attrs.extrapolation_value, ) + + +@register_legalize("relax.image.dilation2d") +def _image_dilation2d(bb: BlockBuilder, call: Call) -> Expr: + """Legalize relax.image.dilation2d to TOPI.""" + layout = call.attrs.data_layout + if layout == "NCHW": + topi_func = topi.image.dilation2d_nchw + elif layout == "NHWC": + topi_func = topi.image.dilation2d_nhwc + else: + raise ValueError(f"dilation2d only supports NCHW or NHWC layout, got {layout}") + + return bb.call_te( + topi_func, + call.args[0], + call.args[1], + stride=call.attrs.strides, + padding=call.attrs.padding, + dilations=call.attrs.dilations, + out_dtype=call.attrs.out_dtype if call.attrs.out_dtype != "" else None, + function_name_hint="dilation2d", + ) diff --git a/python/tvm/relax/transform/legalize_ops/nn.py b/python/tvm/relax/transform/legalize_ops/nn.py index c3c9d317d312..4c660e8d041d 100644 --- a/python/tvm/relax/transform/legalize_ops/nn.py +++ b/python/tvm/relax/transform/legalize_ops/nn.py @@ -790,3 +790,53 @@ def _nn_batch_flatten(bb: BlockBuilder, call: Call) -> Expr: if call.ty.shape is None: return call return bb.call_te(topi.reshape, call.args[0], call.ty.shape.values) + + +@register_legalize("relax.nn.correlation") +def _nn_correlation(bb: BlockBuilder, call: Call) -> Expr: + """Legalize relax.nn.correlation to TOPI.""" + if call.attrs.data_layout != "NCHW": + logging.info( + "TOPI correlation only supports NCHW layout, " + "and thus cannot be legalized by TOPI for other layouts" + ) + return call + return bb.call_te( + topi.nn.correlation_nchw, + call.args[0], + call.args[1], + kernel_size=call.attrs.kernel_size, + max_displacement=call.attrs.max_displacement, + stride1=call.attrs.stride1, + stride2=call.attrs.stride2, + padding=call.attrs.padding, + is_multiply=call.attrs.is_multiply, + function_name_hint="correlation", + ) + + +@register_legalize("relax.nn.lrn") +def _nn_lrn(bb: BlockBuilder, call: Call) -> Expr: + """Legalize relax.nn.lrn (local response normalization) to TOPI.""" + return bb.call_te( + topi.nn.lrn, + call.args[0], + size=call.attrs.size, + axis=call.attrs.axis, + alpha=call.attrs.alpha, + beta=call.attrs.beta, + bias=call.attrs.bias, + function_name_hint="lrn", + ) + + +@register_legalize("relax.nn.dilate") +def _nn_dilate(bb: BlockBuilder, call: Call) -> Expr: + """Legalize relax.nn.dilate to TOPI.""" + return bb.call_te( + topi.nn.dilate, + call.args[0], + strides=call.attrs.strides, + dilation_value=call.attrs.dilation_value, + function_name_hint="dilate", + ) diff --git a/src/relax/op/image/resize.cc b/src/relax/op/image/resize.cc index ae683aee09a5..422c87aaa220 100644 --- a/src/relax/op/image/resize.cc +++ b/src/relax/op/image/resize.cc @@ -142,6 +142,18 @@ void AffineGridAttrs::RegisterReflection() { refl::DefaultValue(true)); } +void Dilation2DAttrs::RegisterReflection() { + namespace refl = tvm::ffi::reflection; + refl::ObjectDef() + .def_ro("strides", &Dilation2DAttrs::strides, "Stride size for the filter window.") + .def_ro("padding", &Dilation2DAttrs::padding, + "Padding size as [pad_top, pad_left, pad_bottom, pad_right].") + .def_ro("dilations", &Dilation2DAttrs::dilations, "Dilation size for the filter.") + .def_ro("data_layout", &Dilation2DAttrs::data_layout, + "Layout of the input data ('NCHW' or 'NHWC').", refl::DefaultValue(ffi::String("NCHW"))) + .def_ro("out_dtype", &Dilation2DAttrs::out_dtype, "Specifies the output data type."); +} + TVM_FFI_STATIC_INIT_BLOCK() { Resize2DAttrs::RegisterReflection(); Resize3DAttrs::RegisterReflection(); @@ -578,5 +590,9 @@ TVM_FFI_STATIC_INIT_BLOCK() { .set_attr(tvm::relax::op_attr::kPurity, true); } +TVM_FFI_STATIC_INIT_BLOCK() { + Dilation2DAttrs::RegisterReflection(); +} + } // namespace relax } // namespace tvm diff --git a/src/relax/op/nn/nn.cc b/src/relax/op/nn/nn.cc index b566cba643e8..fa48deaaaf9f 100644 --- a/src/relax/op/nn/nn.cc +++ b/src/relax/op/nn/nn.cc @@ -177,6 +177,41 @@ void PixelShuffleAttrs::RegisterReflection() { "Scale factor for spatial upsampling."); } +void CorrelationAttrs::RegisterReflection() { + namespace refl = tvm::ffi::reflection; + refl::ObjectDef() + .def_ro("kernel_size", &CorrelationAttrs::kernel_size, + "Kernel size for correlation, must be a positive odd number.") + .def_ro("max_displacement", &CorrelationAttrs::max_displacement, + "Max displacement of correlation.") + .def_ro("stride1", &CorrelationAttrs::stride1, + "Stride for sampling data1 positions.") + .def_ro("stride2", &CorrelationAttrs::stride2, + "Stride for sampling data2 within the neighborhood.") + .def_ro("padding", &CorrelationAttrs::padding, "Padding size for both dimensions.") + .def_ro("is_multiply", &CorrelationAttrs::is_multiply, + "If true, compute product; otherwise compute absolute difference.") + .def_ro("data_layout", &CorrelationAttrs::data_layout, "Layout of the input data."); +} + +void LRNAattrs::RegisterReflection() { + namespace refl = tvm::ffi::reflection; + refl::ObjectDef() + .def_ro("size", &LRNAattrs::size, "Normalization window size.") + .def_ro("axis", &LRNAattrs::axis, "The axis along which normalization is applied.") + .def_ro("alpha", &LRNAattrs::alpha, "The scaling factor.") + .def_ro("beta", &LRNAattrs::beta, "The exponent.") + .def_ro("bias", &LRNAattrs::bias, "Offset to avoid dividing by zero."); +} + +void DilateAttrs::RegisterReflection() { + namespace refl = tvm::ffi::reflection; + refl::ObjectDef() + .def_ro("strides", &DilateAttrs::strides, "Dilation stride for each dimension.") + .def_ro("dilation_value", &DilateAttrs::dilation_value, + "The value used to fill the inserted positions."); +} + TVM_FFI_STATIC_INIT_BLOCK() { SoftmaxAttrs::RegisterReflection(); LeakyReluAttrs::RegisterReflection(); @@ -191,6 +226,9 @@ TVM_FFI_STATIC_INIT_BLOCK() { DropoutAttrs::RegisterReflection(); PadAttrs::RegisterReflection(); PixelShuffleAttrs::RegisterReflection(); + CorrelationAttrs::RegisterReflection(); + LRNAattrs::RegisterReflection(); + DilateAttrs::RegisterReflection(); } /* relax.nn.relu */ From 3d4a15b1f46057180cdd88b977024085b5e5bf87 Mon Sep 17 00:00:00 2001 From: Asthenia0412 Date: Sat, 17 Oct 2026 04:30:32 +0000 Subject: [PATCH 2/4] [FIX][Relax] Register OpDef for new TOPI-backed operators PR #20758 added correlation, lrn, dilate, dilation2d ops with Attribute reflection registration, but missed the corresponding OpDef registration in the C++ OpRegistry. Without OpDef, Python side legalize_ops cannot find the Op via Op.get(), causing 'Operator not registered' errors during import. Add OpDef for: relax.image.dilation2d, relax.nn.correlation, relax.nn.lrn, relax.nn.dilate. --- src/relax/op/image/resize.cc | 6 ++++++ src/relax/op/nn/nn.cc | 22 ++++++++++++++++++++++ 2 files changed, 28 insertions(+) diff --git a/src/relax/op/image/resize.cc b/src/relax/op/image/resize.cc index 422c87aaa220..44404168a2ab 100644 --- a/src/relax/op/image/resize.cc +++ b/src/relax/op/image/resize.cc @@ -592,6 +592,12 @@ TVM_FFI_STATIC_INIT_BLOCK() { TVM_FFI_STATIC_INIT_BLOCK() { Dilation2DAttrs::RegisterReflection(); + + OpDef("relax.image.dilation2d") + .signature(sig::arg("input", "The input tensor."), sig::arg("filter", "The filter tensor."), + sig::call_attrs()) + .set_attr(tvm::relax::op_attr::kMixedPrecisionPolicy, MixedPrecisionPolicyKind::kFollow) + .set_attr(tvm::relax::op_attr::kPurity, true); } } // namespace relax diff --git a/src/relax/op/nn/nn.cc b/src/relax/op/nn/nn.cc index fa48deaaaf9f..8c2a111e09f2 100644 --- a/src/relax/op/nn/nn.cc +++ b/src/relax/op/nn/nn.cc @@ -1516,5 +1516,27 @@ TVM_FFI_STATIC_INIT_BLOCK() { .set_attr(tvm::relax::op_attr::kPurity, true); } +TVM_FFI_STATIC_INIT_BLOCK() { + OpDef("relax.nn.correlation") + .signature(sig::arg("data1", "The first input tensor."), sig::arg("data2", "The second input tensor."), + sig::call_attrs()) + .set_attr(tvm::relax::op_attr::kMixedPrecisionPolicy, MixedPrecisionPolicyKind::kFollow) + .set_attr(tvm::relax::op_attr::kPurity, true); +} + +TVM_FFI_STATIC_INIT_BLOCK() { + OpDef("relax.nn.lrn") + .signature(sig::arg("data", "The input tensor."), sig::call_attrs()) + .set_attr(tvm::relax::op_attr::kMixedPrecisionPolicy, MixedPrecisionPolicyKind::kFollow) + .set_attr(tvm::relax::op_attr::kPurity, true); +} + +TVM_FFI_STATIC_INIT_BLOCK() { + OpDef("relax.nn.dilate") + .signature(sig::arg("data", "The input tensor."), sig::call_attrs()) + .set_attr(tvm::relax::op_attr::kMixedPrecisionPolicy, MixedPrecisionPolicyKind::kFollow) + .set_attr(tvm::relax::op_attr::kPurity, true); +} + } // namespace relax } // namespace tvm From 385113b1461c16caac686a809bc01a9a55bd93b3 Mon Sep 17 00:00:00 2001 From: CI Fix Bot Date: Sat, 17 Oct 2026 10:55:50 +0000 Subject: [PATCH 3/4] [FIX] Fix pre-commit formatting issues: ruff UP038 and clang-format - Fix isinstance() calls to use X | Y union syntax (UP038) - Apply ruff-format to Python files - Apply clang-format to C++ files --- python/tvm/relax/op/image/image.py | 6 ++---- python/tvm/relax/op/nn/nn.py | 2 +- src/relax/op/image/resize.cc | 3 ++- src/relax/op/nn/nn.cc | 7 +++---- 4 files changed, 8 insertions(+), 10 deletions(-) diff --git a/python/tvm/relax/op/image/image.py b/python/tvm/relax/op/image/image.py index abe06ade4010..c0c2367a45d5 100644 --- a/python/tvm/relax/op/image/image.py +++ b/python/tvm/relax/op/image/image.py @@ -389,15 +389,13 @@ def dilation2d( if isinstance(padding, int): padding = (padding, padding, padding, padding) - elif isinstance(padding, (tuple, list)): + elif isinstance(padding, tuple | list): if len(padding) == 2: padding = (padding[0], padding[1], padding[0], padding[1]) elif len(padding) == 4: padding = tuple(padding) else: - raise ValueError( - f"padding must be int, 2-tuple, or 4-tuple, got {len(padding)}-tuple" - ) + raise ValueError(f"padding must be int, 2-tuple, or 4-tuple, got {len(padding)}-tuple") else: raise TypeError(f"padding must be int, tuple, or list, got {type(padding)}") diff --git a/python/tvm/relax/op/nn/nn.py b/python/tvm/relax/op/nn/nn.py index ed8653882f6d..ce864f239e7e 100644 --- a/python/tvm/relax/op/nn/nn.py +++ b/python/tvm/relax/op/nn/nn.py @@ -3205,7 +3205,7 @@ def dilate( result : relax.Expr The dilated output tensor. """ - if not isinstance(strides, (list, tuple)): + if not isinstance(strides, list | tuple): raise TypeError(f"strides must be a list or tuple of ints, got {type(strides)}") strides = [int(s) for s in strides] return _Call( diff --git a/src/relax/op/image/resize.cc b/src/relax/op/image/resize.cc index 44404168a2ab..248cd5cd8175 100644 --- a/src/relax/op/image/resize.cc +++ b/src/relax/op/image/resize.cc @@ -150,7 +150,8 @@ void Dilation2DAttrs::RegisterReflection() { "Padding size as [pad_top, pad_left, pad_bottom, pad_right].") .def_ro("dilations", &Dilation2DAttrs::dilations, "Dilation size for the filter.") .def_ro("data_layout", &Dilation2DAttrs::data_layout, - "Layout of the input data ('NCHW' or 'NHWC').", refl::DefaultValue(ffi::String("NCHW"))) + "Layout of the input data ('NCHW' or 'NHWC').", + refl::DefaultValue(ffi::String("NCHW"))) .def_ro("out_dtype", &Dilation2DAttrs::out_dtype, "Specifies the output data type."); } diff --git a/src/relax/op/nn/nn.cc b/src/relax/op/nn/nn.cc index 8c2a111e09f2..f96b7d398425 100644 --- a/src/relax/op/nn/nn.cc +++ b/src/relax/op/nn/nn.cc @@ -184,8 +184,7 @@ void CorrelationAttrs::RegisterReflection() { "Kernel size for correlation, must be a positive odd number.") .def_ro("max_displacement", &CorrelationAttrs::max_displacement, "Max displacement of correlation.") - .def_ro("stride1", &CorrelationAttrs::stride1, - "Stride for sampling data1 positions.") + .def_ro("stride1", &CorrelationAttrs::stride1, "Stride for sampling data1 positions.") .def_ro("stride2", &CorrelationAttrs::stride2, "Stride for sampling data2 within the neighborhood.") .def_ro("padding", &CorrelationAttrs::padding, "Padding size for both dimensions.") @@ -1518,8 +1517,8 @@ TVM_FFI_STATIC_INIT_BLOCK() { TVM_FFI_STATIC_INIT_BLOCK() { OpDef("relax.nn.correlation") - .signature(sig::arg("data1", "The first input tensor."), sig::arg("data2", "The second input tensor."), - sig::call_attrs()) + .signature(sig::arg("data1", "The first input tensor."), + sig::arg("data2", "The second input tensor."), sig::call_attrs()) .set_attr(tvm::relax::op_attr::kMixedPrecisionPolicy, MixedPrecisionPolicyKind::kFollow) .set_attr(tvm::relax::op_attr::kPurity, true); } From 0c03acbd36244c2c2e20a34247e5d6db65819786 Mon Sep 17 00:00:00 2001 From: CI Fix Bot Date: Sat, 17 Oct 2026 11:42:46 +0000 Subject: [PATCH 4/4] [FIX] Use tvm_ffi.dtype instead of DataType to fix NameError - Remove incorrect 'from tvm_tirx import DataType' imports from image.py and nn.py - Replace all 'DataType' type annotations with 'tvm_ffi.dtype' - The original PR mistakenly replaced tvm_ffi.dtype with DataType, which caused NameError at module import time since DataType is not available in the runtime namespace. --- python/tvm/relax/op/image/image.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/tvm/relax/op/image/image.py b/python/tvm/relax/op/image/image.py index c0c2367a45d5..817c0d5de56c 100644 --- a/python/tvm/relax/op/image/image.py +++ b/python/tvm/relax/op/image/image.py @@ -331,7 +331,7 @@ def dilation2d( padding: int | tuple[int, int, int, int] | str = 0, dilations: int | tuple[int, int] = (1, 1), data_layout: str = "NCHW", - out_dtype: str | DataType | None = None, + out_dtype: str | tvm_ffi.dtype | None = None, *, ty=None, loc: Location = UNKNOWN_LOC, @@ -374,7 +374,7 @@ def dilation2d( data_layout : str Layout of the input data. Supported: "NCHW", "NHWC". - out_dtype : Optional[str | DataType] + out_dtype : Optional[str | tvm_ffi.dtype] Specifies the output data type. If None, uses the input dtype. Returns