I get an exception when printing an IR from a pytorch imported model with keep_params_as_input=True, the error doesn’t happen with the default False value.
I am running an updated version of unity - as of yesterday.
Happy to file an issue if needed.
Here is a minimal example:
class TorchNN(nn.Module):
def __init__(self):
super().__init__()
self.flatten = nn.Flatten()
self.linear_relu_stack = nn.Sequential(
nn.Linear(28, 512, bias=False),
)
def forward(self, x):
x = self.flatten(x)
logits = self.linear_relu_stack(x)
return logits
network = TorchNN()
traced_network = fx.symbolic_trace(network)
input_shape = (1, 28, 28)
input_dtype = "float"
input_info = [(input_shape, input_dtype)]
mod = relax.frontend.torch.from_fx(traced_network, input_info, keep_params_as_input=True)
mod.show()
Error:
[11:07:25] /tvm/include/tvm/script/printer/ir_docsifier_functor.h:72: Warning: ObjectFunctor calls un-registered function on type: runtime.NDArray (token: relax). ObjectType: runtime.NDArray. Object: runtime.NDArray(0x600003a38f00)
---------------------------------------------------------------------------
TVMError Traceback (most recent call last)
Cell In[254], line 1
----> 1 irmod.show()
File /tvm/python/tvm/runtime/script_printer.py:328, in Scriptable.show(self, style, black_format, name, show_meta, ir_prefix, tir_prefix, relax_prefix, module_alias, buffer_dtype, int_dtype, float_dtype, verbose_expr, indent_spaces, print_line_numbers, num_context_lines, syntax_sugar, path_to_underline, path_to_annotate, obj_to_underline, obj_to_annotate)
276 """A sugar for print highlighted TVM script.
277
278 Parameters
(...)
321 Object to be annotated
322 """
323 from tvm.script.highlight import ( # pylint: disable=import-outside-toplevel
324 cprint,
325 )
327 cprint(
--> 328 self.script(
329 name=name,
330 show_meta=show_meta,
331 ir_prefix=ir_prefix,
332 tir_prefix=tir_prefix,
333 relax_prefix=relax_prefix,
334 module_alias=module_alias,
335 buffer_dtype=buffer_dtype,
336 int_dtype=int_dtype,
337 float_dtype=float_dtype,
338 verbose_expr=verbose_expr,
339 indent_spaces=indent_spaces,
340 print_line_numbers=print_line_numbers,
341 num_context_lines=num_context_lines,
342 syntax_sugar=syntax_sugar,
343 path_to_underline=path_to_underline,
344 path_to_annotate=path_to_annotate,
345 obj_to_underline=obj_to_underline,
346 obj_to_annotate=obj_to_annotate,
347 ),
348 style=style,
349 black_format=black_format,
350 )
File /tvm/python/tvm/runtime/script_printer.py:182, in Scriptable.script(self, name, show_meta, ir_prefix, tir_prefix, relax_prefix, module_alias, buffer_dtype, int_dtype, float_dtype, verbose_expr, indent_spaces, print_line_numbers, num_context_lines, syntax_sugar, path_to_underline, path_to_annotate, obj_to_underline, obj_to_annotate)
113 def script(
114 self,
115 *,
(...)
133 obj_to_annotate: Optional[Dict[Object, str]] = None,
134 ) -> str:
135 """Print TVM IR into TVMScript text format
136
137 Parameters
(...)
180 The TVM Script of the given TVM IR
181 """
--> 182 return _script(
183 self,
184 PrinterConfig(
185 name=name,
186 show_meta=show_meta,
187 ir_prefix=ir_prefix,
188 tir_prefix=tir_prefix,
189 relax_prefix=relax_prefix,
190 module_alias=module_alias,
191 buffer_dtype=buffer_dtype,
192 int_dtype=int_dtype,
193 float_dtype=float_dtype,
194 verbose_expr=verbose_expr,
195 indent_spaces=indent_spaces,
196 print_line_numbers=print_line_numbers,
197 num_context_lines=num_context_lines,
198 syntax_sugar=syntax_sugar,
199 path_to_underline=path_to_underline,
200 path_to_annotate=path_to_annotate,
201 obj_to_underline=obj_to_underline,
202 obj_to_annotate=obj_to_annotate,
203 ),
204 )
File /tvm/python/tvm/runtime/script_printer.py:102, in _script(obj, config)
101 def _script(obj: Object, config: PrinterConfig) -> str:
--> 102 return _ffi_node_api.TVMScriptPrinterScript(obj, config)
File /tvm/python/tvm/_ffi/_ctypes/packed_func.py:238, in PackedFuncBase.__call__(self, *args)
226 ret_tcode = ctypes.c_int()
227 if (
228 _LIB.TVMFuncCall(
229 self.handle,
(...)
236 != 0
237 ):
--> 238 raise get_last_ffi_error()
239 _ = temp_args
240 _ = args
TVMError: Traceback (most recent call last):
File "/tvm/include/tvm/script/printer/ir_docsifier_functor.h", line 75
TVMError:
---------------------------------------------------------------
An error occurred during the execution of TVM.
For more information, please see: https://tvm.apache.org/docs/errors.html
---------------------------------------------------------------
Check failed: (false) is false: ObjectFunctor calls un-registered function on type: runtime.NDArray (token: relax). ObjectType: runtime.NDArray. Object: runtime.NDArray(0x600003a38f00)