Method __call__

imperative/python/megengine/core/autodiff/grad.py:180–215

(self, *args)

        return self._default_rule(*args), self.backward

    def __call__(self, *args):
        from ...tensor import Tensor

        for arg in args:
            if not isinstance(arg, Tensor):
                raise TypeError(
                    "op Function expect type Tensor as inputs, got {}".format(type(arg))
                )

        grad_key = core2.get_grad_key(args)
        if grad_key is None:
            return self._default_rule(*args)

        grad = Grad.key2grad[grad_key]
        group = [ref() for ref in grad._group]

        origin_args = [Tensor(arg) for arg in args]

        for grad in group:
            grad.suppress()
        outputs, backward = self._grad_rule(*args)
        for grad in reversed(group):
            grad.resume()

        def normalized_backward(*output_grads):
            input_grads = backward(*output_grads)
            if isinstance(input_grads, Tensor) or input_grads is None:
                input_grads = (input_grads,)
            return input_grads

        if self.__single_output:
            outputs = (outputs,)
        outputs = core2.set_grad(normalized_backward, origin_args, outputs)
        if self.__single_output:
            (outputs,) = outputs
        return outputs

    def __getstate__(self):
        return self.__dict__