Method flatten_inputs

imperative/python/megengine/jit/tracing.py:244–274

(self, *args, **kwargs)

        return self._trace.without_host

    def flatten_inputs(self, *args, **kwargs):
        from ..traced_module.pytree import tree_flatten, SUPPORTED_LEAF_CLS
        from ..module import Module
        from ..optimizer import Optimizer

        if Optimizer not in SUPPORTED_LEAF_CLS:
            SUPPORTED_LEAF_CLS.append(Optimizer)

        tensor_args = []
        modules = []
        fargs, _ = tree_flatten((args, kwargs))
        for a in fargs:
            if isinstance(a, RawTensor):
                tensor_args.append(a)
            elif isinstance(a, Module):
                modules.append(a)
            elif isinstance(a, Optimizer):
                states = a._state.values()
                for s in states:
                    assert isinstance(s, dict)
                    tensor_args.extend(list(s.values()))
        for m in modules:
            tensor_args.extend(list(m.parameters()))
        grads = []
        for t in tensor_args:
            if t.grad is not None:
                grads.append(t.grad)
        for m in modules:
            tensor_args.extend(list(m.buffers()))
        tensor_args.extend(grads)
        return tensor_args

    def compile(self):
        raise NotImplementedError