Annotate a Relax function. Parameters ---------- name : str, optional The name of the function params : tvm.relax.Var | Tuple | List[tvm.relax.Var], optional The parameters of the function. If params is None, it means deferring in
(
self,
name: str,
params: Var | Tuple | list[Var] | None = None,
attrs: dict[str, Object] | None = None,
pure: bool = True,
private: bool = False,
)
| 211 | raise RuntimeError("emit_func_output must be called in a relax function.") |
| 212 | |
| 213 | def function( |
| 214 | self, |
| 215 | name: str, |
| 216 | params: Var | Tuple | list[Var] | None = None, |
| 217 | attrs: dict[str, Object] | None = None, |
| 218 | pure: bool = True, |
| 219 | private: bool = False, |
| 220 | ) -> FunctionScope: |
| 221 | """Annotate a Relax function. |
| 222 | |
| 223 | Parameters |
| 224 | ---------- |
| 225 | name : str, optional |
| 226 | The name of the function |
| 227 | |
| 228 | params : tvm.relax.Var | Tuple | List[tvm.relax.Var], optional |
| 229 | The parameters of the function. |
| 230 | If params is None, it means deferring initialization of function parameters |
| 231 | until emit_func_output. |
| 232 | |
| 233 | attrs : Dict[str, Object], optional |
| 234 | The function attrs |
| 235 | |
| 236 | pure : bool, optional |
| 237 | Whether the function is annotated as pure. |
| 238 | |
| 239 | private : bool, optional |
| 240 | Whether the function is annotated as private. |
| 241 | If the function is private, it will not have a global symbol attribute. |
| 242 | If it is not private and not an inner function, then it will have |
| 243 | a global symbol attribute (mapped to the function's name) |
| 244 | |
| 245 | Returns |
| 246 | ------- |
| 247 | ret: FunctionScope |
| 248 | A FunctionScope for building a Relax function node. |
| 249 | """ |
| 250 | if isinstance(params, rx.Var): |
| 251 | params = [params] |
| 252 | elif isinstance(params, list | tuple): |
| 253 | for param in params: |
| 254 | if not isinstance(param, rx.Var): |
| 255 | raise TypeError( |
| 256 | f"each element of function parameters must be of type tvm.relax.Var,\ |
| 257 | but got: {type(param)}" |
| 258 | ) |
| 259 | if attrs is None: |
| 260 | attrs = {} |
| 261 | # The block builder does not permit nesting functions, per above comment, |
| 262 | # so no further check should be needed |
| 263 | if not private: |
| 264 | attrs["global_symbol"] = name |
| 265 | |
| 266 | return FunctionScope(self, name, params, attrs, is_pure=pure) |
| 267 | |
| 268 | def testing_scope(self, def_vars: list[tirx.Var]) -> TestingScope: |
| 269 | """Start a scope for unit-testing purposes. |