r"""``CallMethod`` represents a call to the ``__call__`` method of ``Module`` or a method of ``Tensor``. Args: node: the Node to be called. method: the method name. Default: "__call__"
| 420 | |
| 421 | # expr: outputs = inputs[0].__call__(*inputs[1:]) |
| 422 | class CallMethod(Expr): |
| 423 | r"""``CallMethod`` represents a call to the ``__call__`` method of ``Module`` or a method of ``Tensor``. |
| 424 | |
| 425 | Args: |
| 426 | node: the Node to be called. |
| 427 | method: the method name. |
| 428 | Default: "__call__" |
| 429 | """ |
| 430 | |
| 431 | def __init__(self, node, method="__call__"): |
| 432 | super().__init__() |
| 433 | if isinstance(node, type): |
| 434 | assert issubclass(node, Tensor) |
| 435 | cls = Parameter if issubclass(node, Parameter) else Tensor |
| 436 | |
| 437 | self.inputs = [] |
| 438 | self.const_val = [(0, cls)] |
| 439 | else: |
| 440 | assert isinstance(node, (TensorNode, ModuleNode)) |
| 441 | node.users.append(self) |
| 442 | self.inputs = [ |
| 443 | node, |
| 444 | ] |
| 445 | self.const_val = [] |
| 446 | self.arg_def = tree_flatten(((node,), {}))[1] |
| 447 | self.method = method |
| 448 | |
| 449 | @classmethod |
| 450 | def make(cls, *args, **kwargs): |
| 451 | assert active_module_tracer() is not None |
| 452 | expr = cls(*args, **kwargs) |
| 453 | active_module_tracer().current_scope()._insert(expr) |
| 454 | return expr |
| 455 | |
| 456 | @property |
| 457 | def graph(self): |
| 458 | if isinstance(self.inputs[0], ModuleNode): |
| 459 | m_node = self.inputs[0] |
| 460 | if ( |
| 461 | hasattr(m_node.owner, "argdef_graph_map") |
| 462 | and m_node.owner.argdef_graph_map |
| 463 | ): |
| 464 | assert self.arg_def in m_node.owner.argdef_graph_map |
| 465 | return m_node.owner.argdef_graph_map[self.arg_def] |
| 466 | return None |
| 467 | |
| 468 | def interpret(self, *inputs): |
| 469 | args, kwargs = self.unflatten_args(inputs) |
| 470 | obj = args[0] |
| 471 | meth = getattr(obj, self.method) |
| 472 | if inspect.ismethod(meth): |
| 473 | args = args[1:] |
| 474 | outputs = getattr(obj, self.method)(*args, **kwargs) |
| 475 | if self.method == "__setitem__": |
| 476 | outputs = obj |
| 477 | if outputs is None: |
| 478 | return outputs |
| 479 | outputs, _ = tree_flatten(outputs, is_leaf=lambda x: isinstance(x, RawTensor)) |
no outgoing calls