chore: import upstream snapshot with attribution

python/paddle/distributed/auto_parallel/static/graph.py

@@ -0,0 +1,186 @@
+#   Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+
+from collections import OrderedDict
+
+
+class Node:
+    def __init__(self, id, **attrs):
+        # Each node must has a unique id
+        self._id = id
+        # Attributes for Node
+        self._attrs = {}
+        self._attrs.update(attrs)
+
+    @property
+    def id(self):
+        return self._id
+
+    @property
+    def attrs(self):
+        return self._attrs
+
+    def __getitem__(self, attr_name):
+        return self._attrs[attr_name]
+
+    def __setitem__(self, attr_name, attr_value):
+        self._attrs[attr_name] = attr_value
+
+    def __contains__(self, attr_name):
+        try:
+            return attr_name in self._attrs
+        except TypeError:
+            return False
+
+    def __str__(self):
+        str = f"(id: {self.id}, attrs: {self.attrs})"
+        return str
+
+
+class Edge:
+    def __init__(self, src_id, tgt_id, **attrs):
+        # The id of source node in an Edge
+        self._src_id = src_id
+        # The id of target node in an Edge
+        self._tgt_id = tgt_id
+        # Attributes for Edge
+        self._attrs = {}
+        self._attrs.update(attrs)
+
+    @property
+    def src_id(self):
+        return self._src_id
+
+    @property
+    def tgt_id(self):
+        return self._tgt_id
+
+    @property
+    def attrs(self):
+        return self._attrs
+
+    def __getitem__(self, attr_name):
+        return self._attrs[attr_name]
+
+    def __setitem__(self, attr_name, attr_value):
+        self._attrs[attr_name] = attr_value
+
+    def __contains__(self, attr_name):
+        try:
+            return attr_name in self._attrs
+        except TypeError:
+            return False
+
+    def __str__(self):
+        str = ""
+        str += f"(src_id: {self.src_id}, tgt_id: {self.tgt_id}, attrs: {self._attrs})"
+        return str
+
+
+class Graph:
+    def __init__(self, **attrs):
+        # _nodes is dict for storing the nodes of the graph.
+        # The key of this dict is the node id.
+        self._nodes = {}
+        # _adjs is a dict of dict for storing the adjacency of the graph.
+        # The key of the outer dict is the node id of the source node and
+        # the key of the inner dict is the node id of the target node.
+        self._adjs = {}
+        # Attributes for Graph
+        self._attrs = {}
+        self._attrs.update(attrs)
+        self._reverse_adjs = {}
+        self._attr_to_nodes = {}
+
+    @property
+    def nodes(self):
+        return self._nodes
+
+    @property
+    def attrs(self):
+        return self._attrs
+
+    @property
+    def adjs(self):
+        return self._adjs
+
+    def add_node(self, node_id, **attrs):
+        if node_id is None:
+            raise ValueError("None cannot be a node")
+        if node_id not in self._nodes:
+            node = Node(node_id, **attrs)
+            self._nodes[node_id] = node
+            self._adjs[node_id] = {}
+            self._reverse_adjs[node_id] = []
+        else:
+            self._nodes[node_id].attrs.update(attrs)
+
+        return self._nodes[node_id]
+
+    def add_edge(self, src_id, tgt_id, **attrs):
+        # add nodes
+        if src_id is None:
+            raise ValueError("None cannot be a node")
+        if tgt_id is None:
+            raise ValueError("None cannot be a node")
+        if src_id not in self._nodes:
+            src_node = Node(src_id)
+            self._nodes[src_id] = src_node
+            # for one tensor to multiple ops
+            self._adjs[src_id] = OrderedDict()
+            self._reverse_adjs[src_id] = []
+        if tgt_id not in self._nodes:
+            tgt_node = Node(tgt_id)
+            self._nodes[tgt_id] = tgt_node
+            # for one tensor to multiple ops
+            self._adjs[tgt_id] = OrderedDict()
+            self._reverse_adjs[tgt_id] = []
+        # add the edge
+        edge = Edge(src_id, tgt_id, **attrs)
+        self._adjs[src_id][tgt_id] = edge
+
+        # add the reverse adj
+        self._reverse_adjs[tgt_id].append(self.nodes[src_id])
+        return edge
+
+    def __len__(self):
+        return len(self._nodes)
+
+    def __iter__(self):
+        return iter(self._nodes.values())
+
+    def __getitem__(self, node_id):
+        # Return the adjacency of a node
+        return self._adjs[node_id]
+
+    def __contains__(self, node_id):
+        # Check whether a node in the graph
+        try:
+            return node_id in self._nodes
+        except TypeError:
+            return False
+
+    def __str__(self):
+        str = ""
+        str += "**************Nodes**************\n"
+        for node_id in self.nodes:
+            str += f"{self.nodes[node_id]}\n"
+
+        str += "**************Edges**************\n"
+        for src_id in self.adjs:
+            str += f"--------------{src_id}--------------\n"
+            for idx, tgt_id in enumerate(self.adjs[src_id]):
+                str += f"{self.adjs[src_id][tgt_id]}\n"
+
+        return str