mbs-octoml commented on a change in pull request #9038:
URL: https://github.com/apache/tvm/pull/9038#discussion_r717072043
##########
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
+def function_on_device(function, param_devices, result_device):
Review comment:
I ended up fully commenting this one in a prior PR.
--
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]
