Method __map_io

python/oneflow/nn/graph/proxy.py:349–391

(self, io_type, func, func_desc, *args, **kwargs)

                    yield m

    def __map_io(self, io_type, func, func_desc, *args, **kwargs):
        assert isinstance(func_desc, str)
        assert io_type in ("input", "output")
        mapped_args = []

        def map_tensor(item):
            assert isinstance(item, Tensor)
            return func(item)

        args_tree = ArgsTree(
            (args, kwargs),
            True,
            "_"
            + self.to(GraphModule).name_prefix
            + self.to(GraphModule).name
            + "_"
            + io_type,
            None,
        )

        def leaf_node_fn(leaf_node):
            arg = leaf_node.value()
            name = leaf_node.prefix() + "_" + leaf_node.name()
            is_tensor, repr_str = self.__io_tensor_check_and_gen(arg, io_type, name)
            if is_tensor:
                self.__print(
                    0,
                    1,
                    f"{repr_str} is a Tensor, {func_desc} transformation has been done.",
                )
                return map_tensor(arg)
            else:
                self.__print(
                    0,
                    0,
                    f"{repr_str} is not a Tensor, {func_desc} transformation will be ignored.",
                )
                return arg

        out = args_tree.map_leaf(leaf_node_fn)
        mapped_args = out[0]
        mapped_kwargs = out[1]
        return mapped_args, mapped_kwargs

    def __io_tensor_check_and_gen(self, item, io_type, name):
        assert io_type in ("input", "output")