electriclilies commented on a change in pull request #10352:
URL: https://github.com/apache/tvm/pull/10352#discussion_r813441337
##########
File path: src/relay/ir/expr_functor.cc
##########
@@ -476,41 +476,31 @@ Expr Bind(const Expr& expr, const tvm::Map<Var, Expr>&
args_map) {
if (const FunctionNode* func = expr.as<FunctionNode>()) {
Expr new_body = ExprBinder(args_map).VisitExpr(func->body);
Array<Var> new_params;
- std::vector<VirtualDevice> new_param_virtual_devices;
+ bool params_unchanged = true;
for (size_t i = 0; i < func->params.size(); ++i) {
if (!args_map.count(func->params[i])) {
new_params.push_back(func->params[i]);
-
new_param_virtual_devices.push_back(GetFunctionParamVirtualDevice(func, i));
+ } else if (const auto var = args_map[func->params[i]].as<VarNode>()) {
+ // If we're mapping a variable to a variable and not a normal expr,
then we want to
+ // put the substitution in the new parameters.
+ params_unchanged = false;
+ new_params.push_back(GetRef<Var>(var));
}
}
- if (new_body.same_as(func->body) && new_params.size() ==
func->params.size()) {
+ if (new_body.same_as(func->body) && new_params.size() ==
func->params.size() &&
+ params_unchanged) {
return expr;
}
auto ret =
Function(new_params, new_body, func->ret_type, func->type_params,
func->attrs, func->span);
- ret =
- MaybeFunctionOnDevice(ret, new_param_virtual_devices,
GetFunctionResultVirtualDevice(func));
- std::unordered_set<Var, ObjectPtrHash, ObjectPtrEqual> set;
Review comment:
Yeah I think that makes more sense, I'll revert the changes in Bind and
use ExprBinder instead
--
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]
