Wrapper for Relax function with automatic tensor conversion.
(*args, **kwargs)
| 167 | |
| 168 | def _create_relax_wrapper(name): |
| 169 | def wrapper(*args, **kwargs): |
| 170 | """Wrapper for Relax function with automatic tensor conversion.""" |
| 171 | if hasattr(self.ir_mod, "pyfuncs") and name in self.ir_mod.pyfuncs: |
| 172 | return self.ir_mod.pyfuncs[name](*args, **kwargs) |
| 173 | |
| 174 | if self.relax_vm is not None: |
| 175 | converted_args = self._convert_pytorch_to_tvm(list(args)) |
| 176 | converted_kwargs = { |
| 177 | k: self._convert_pytorch_to_tvm(v) for k, v in kwargs.items() |
| 178 | } |
| 179 | result = self.relax_vm[name](*converted_args, **converted_kwargs) |
| 180 | return self._convert_tvm_to_pytorch(result) |
| 181 | |
| 182 | raise RuntimeError( |
| 183 | f"Neither converted Python function nor Relax VM available for {name}" |
| 184 | ) |
| 185 | |
| 186 | wrapper.__name__ = name |
| 187 | wrapper.__doc__ = f"Wrapped Relax function: {name}" |
nothing calls this directly
no test coverage detected