tlopex commented on code in PR #18313:
URL: https://github.com/apache/tvm/pull/18313#discussion_r2361432743
##########
python/tvm/script/ir_builder/relax/ir.py:
##########
@@ -451,6 +453,55 @@ def call_packed(
return Call(op, args, attrs=attrs, sinfo_args=sinfo_args)
+@args_converter.auto
+def call_py_func(
+ py_func_name: py_str,
+ *args: Expr,
+ out_sinfo: Union[StructInfo, List[StructInfo]],
+) -> Call:
+ """Create a relax Call, which calls a Python function.
+
+ Parameters
+ ----------
+ py_func_name: str
+ The name of the Python function to call. This should correspond to a
function
+ in the IRModule's pyfuncs attribute.
+ *args : Expr
+ The arguments.
+ out_sinfo: Union[StructInfo, List[StructInfo]]
+ The structure info of the call_py_func output.
+ It should be a single or a list of TensorStructInfo. Each one denotes
the
+ structure info of a returned tensor.
+
+ Returns
+ -------
+ call: Call
+ The created Relax Call for call_py_func operator.
+ """
+ if isinstance(out_sinfo, py_tuple): # type: ignore
+ out_sinfo = list(out_sinfo)
+ elif not isinstance(out_sinfo, list):
+ out_sinfo = [out_sinfo]
+
+ out_sinfo = [
+ (
+ sinfo()
+ if callable(sinfo)
+ else sinfo.asobject()
+ if isinstance(sinfo, ObjectConvertible)
+ else sinfo
+ )
+ for sinfo in out_sinfo
+ ]
+
+ # Convert string to StringImm
+ try:
+ func_name_imm = StringImm(py_func_name) if hasattr(py_func_name,
"strip") else py_func_name
Review Comment:
Okay, I think we can use `isinstance(py_func_name, str)`
--
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]
