-
Notifications
You must be signed in to change notification settings - Fork 3.8k
[Relay][VirtualDevice] Expose WithFields to Python to do proper copy in ExprMutator #11882
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Conversation
mbs-octoml
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM, thanks and sorry for the glitch. Just some nits to make the API more convenient as per the c++ API.
python/tvm/relay/expr_functor.py
Outdated
| fn, | ||
| list(new_params), | ||
| new_body, | ||
| fn.ret_type, |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
python/tvm/relay/function.py
Outdated
|
|
||
|
|
||
| @tvm._ffi.register_func("relay.FunctionWithFields") | ||
| def FunctionWithFields(function, params, body, ret_type, ty_params, attrs, virtual_device, span): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think defaulting all params after function to None will mimic the c++ version and support the above suggestion.
src/relay/ir/function.cc
Outdated
| 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, |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Preserve the Optional<...> on all params after function to support the above suggestions.
|
@tvm-bot rerun |
1921f44 to
ad5dcd7
Compare
…in ExprMutator (apache#11882) * [Relay][VirtualDevice] Expose WithFields to Python to do proper copy in ExprMutator * [Relay] give FunctionWithFields optional arguments * [lint] fix wrong line length * [lint] missing newline * [doc] add doc string to FunctionWithFields
…in ExprMutator (apache#11882) * [Relay][VirtualDevice] Expose WithFields to Python to do proper copy in ExprMutator * [Relay] give FunctionWithFields optional arguments * [lint] fix wrong line length * [lint] missing newline * [doc] add doc string to FunctionWithFields
…in ExprMutator (apache#11882) * [Relay][VirtualDevice] Expose WithFields to Python to do proper copy in ExprMutator * [Relay] give FunctionWithFields optional arguments * [lint] fix wrong line length * [lint] missing newline * [doc] add doc string to FunctionWithFields
As discussed in https://discuss.tvm.apache.org/t/relay-function-virtual-device-property/12958
@mbs-octoml