mbs-octoml commented on code in PR #11882:
URL: https://github.com/apache/tvm/pull/11882#discussion_r907604246
##########
python/tvm/relay/function.py:
##########
@@ -63,3 +63,10 @@ def __call__(self, *args):
Arguments.
"""
return Call(self, args, None, None)
+
+
+@tvm._ffi.register_func("relay.FunctionWithFields")
+def FunctionWithFields(function, params, body, ret_type, ty_params, attrs,
virtual_device, span):
Review Comment:
I think defaulting all params after function to None will mimic the c++
version and support the above suggestion.
##########
src/relay/ir/function.cc:
##########
@@ -59,6 +59,12 @@ TVM_REGISTER_GLOBAL("relay.ir.Function")
tvm::Array<TypeVar> ty_params, tvm::DictAttrs attrs) {
return Function(params, body, ret_type, ty_params, attrs);
});
+TVM_REGISTER_GLOBAL("relay.ir.FunctionWithFields")
+ .set_body_typed([](Function function, tvm::Array<Var> params, Expr body,
Type ret_type,
Review Comment:
Preserve the Optional<...> on all params after function to support the above
suggestions.
##########
python/tvm/relay/expr_functor.py:
##########
@@ -202,7 +202,16 @@ class ExprMutator(ExprFunctor):
def visit_function(self, fn):
new_params = [self.visit(x) for x in fn.params]
new_body = self.visit(fn.body)
- return Function(list(new_params), new_body, fn.ret_type,
fn.type_params, fn.attrs)
+ return FunctionWithFields(
+ fn,
+ list(new_params),
+ new_body,
+ fn.ret_type,
Review Comment:
You can elide all the fn.* args since WithFields is designed to make it easy
to preserve everything that's unchanged from the original expression.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
