This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch fix-tvm-ffi-bump-relax-slots in repository https://gitbox.apache.org/repos/asf/tvm.git
commit dfde900cecf90130f9b07665976b85aec867f31d Author: tqchen <[email protected]> AuthorDate: Fri Jul 17 10:29:36 2026 +0000 [IR] Preserve functions across with_attr updates --- python/tvm/ir/function.py | 12 +++++++----- tests/python/relax/test_expr.py | 11 ++++++----- 2 files changed, 13 insertions(+), 10 deletions(-) diff --git a/python/tvm/ir/function.py b/python/tvm/ir/function.py index 76699dcab3..abf7d37372 100644 --- a/python/tvm/ir/function.py +++ b/python/tvm/ir/function.py @@ -62,17 +62,19 @@ class BaseFunc(Expr): func : BaseFunc A new copy of the function """ - # make sure we first copy so that we can safely do copy on write - # for multiple updates. - res = _ffi_api.BaseFuncCopy(self) + # Pass an lvalue so that the RValueRef argument takes its own strong + # reference. tvm-ffi ties a C++ object to one canonical Python wrapper, + # so BaseFuncCopy(self) may return self; moving that wrapper would also + # invalidate the caller's original function. + res = self if isinstance(attr_key_or_dict, dict): for key, val in attr_key_or_dict.items(): - res = _ffi_api.BaseFuncWithAttr(res._move(), key, tvm.runtime.convert(val)) + res = _ffi_api.BaseFuncWithAttr(res, key, tvm.runtime.convert(val)) return res return _ffi_api.BaseFuncWithAttr( - res._move(), attr_key_or_dict, tvm.runtime.convert(attr_value) + res, attr_key_or_dict, tvm.runtime.convert(attr_value) ) def with_attrs(self, attr_map: DictAttrs | dict[str, Object]) -> "BaseFunc": diff --git a/tests/python/relax/test_expr.py b/tests/python/relax/test_expr.py index 52b4e9ba16..a84032ef5b 100644 --- a/tests/python/relax/test_expr.py +++ b/tests/python/relax/test_expr.py @@ -235,11 +235,12 @@ def test_func(): seqe = rx.SeqExpr(blocks, x) ret_ty = R.Tensor(dtype="float32", ndim=-1) func = rx.Function([x], seqe, ret_ty) - func = func.with_attr("global_symbol", "func") - assert func.params[0] == x - assert func.body == seqe - assert func.ret_ty == ret_ty - assert func.attrs["global_symbol"] == "func" + with_attr = func.with_attr("global_symbol", "func") + assert "global_symbol" not in func.attrs + assert with_attr.params[0] == x + assert with_attr.body == seqe + assert with_attr.ret_ty == ret_ty + assert with_attr.attrs["global_symbol"] == "func" def test_shape_of():
