mbs-octoml commented on a change in pull request #9077:
URL: https://github.com/apache/tvm/pull/9077#discussion_r715042461
##########
File path: python/tvm/relay/op/annotation/annotation.py
##########
@@ -33,21 +43,26 @@ def on_device(data, device):
device : Union[:py:class:`Device`, str]
The device type to annotate.
+ is_fixed : bool
+ If true, annotation does not imply a device_copy may be inserted.
+ (This parameter is used internally by the compiler and unit tests and
+ should not need to be set in user programs.)
+
Returns
-------
result : tvm.relay.Expr
The annotated expression.
"""
- if isinstance(device, _Device):
- device = device.device_type
- elif isinstance(device, str):
- device = _nd.device(device).device_type
- else:
- raise ValueError(
- "device is expected to be the type of Device or "
- "str, but received %s" % (type(device))
- )
- return _make.on_device(data, device)
+ return _make.on_device(data, _device_to_int(device), is_fixed)
+
+
+# for testing only
Review comment:
Woops, I changed my mind on this and think it's a legit user-visible
annotation, but forgot to update the comment. Done.
--
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]
