Adjust import statements

mlcommons · May 9, 2024 · 90d5c12 · 90d5c12
1 parent 91803ce
commit 90d5c12
Show file tree

Hide file tree

Showing 5 changed files with 22 additions and 23 deletions.
diff --git a/src/python/converter/pytorch2chakra_converter.py b/src/python/converter/pytorch2chakra_converter.py
@@ -4,7 +4,7 @@
 import logging
 from typing import Dict, List, Optional, Set, Tuple
 
-from chakra.et_def.et_def_pb2 import (
+from ....schema.protobuf.et_def_pb2 import (
     ALL_GATHER,
     ALL_REDUCE,
     ALL_TO_ALL,
@@ -14,17 +14,16 @@
     REDUCE_SCATTER,
     GlobalMetadata,
 )
-from chakra.et_def.et_def_pb2 import (
+from ....schema.protobuf.et_def_pb2 import (
     AttributeProto as ChakraAttr,
 )
-from chakra.et_def.et_def_pb2 import (
+from ....schema.protobuf.et_def_pb2 import (
     Node as ChakraNode,
 )
-from chakra.et_def.et_def_pb2 import (
+from ....schema.protobuf.et_def_pb2 import (
     NodeType as ChakraNodeType,
 )
-from chakra.third_party.utils.protolib import encodeMessage as encode_message
-
+from ..third_party.protolib import encodeMessage as encode_message
 from .pytorch_node import PyTorchNode, PyTorchNodeType
 
 

diff --git a/src/python/converter/text2chakra_converter.py b/src/python/converter/text2chakra_converter.py
@@ -4,7 +4,7 @@
 from io import TextIOWrapper
 from typing import Any, List
 
-from chakra.et_def.et_def_pb2 import (
+from ....schema.protobuf.et_def_pb2 import (
     ALL_GATHER,
     ALL_REDUCE,
     ALL_TO_ALL,
@@ -14,11 +14,9 @@
     GlobalMetadata,
     Node,
     NodeType,
-)
-from chakra.et_def.et_def_pb2 import (
     AttributeProto as ChakraAttr,
 )
-from chakra.third_party.utils.protolib import encodeMessage as encode_message
+from ..third_party.protolib import encodeMessage as encode_message
 
 
 class Layer:

diff --git a/src/python/generator/generator.py b/src/python/generator/generator.py
@@ -2,7 +2,7 @@
 
 import argparse
 
-from chakra.et_def.et_def_pb2 import (
+from ....schema.protobuf.et_def_pb2 import (
     ALL_GATHER,
     ALL_REDUCE,
     ALL_TO_ALL,
@@ -29,16 +29,16 @@
     Uint32List,
     Uint64List,
 )
-from chakra.et_def.et_def_pb2 import (
+from ....schema.protobuf.et_def_pb2 import (
     AttributeProto as ChakraAttr,
 )
-from chakra.et_def.et_def_pb2 import (
+from ....schema.protobuf.et_def_pb2 import (
     Node as ChakraNode,
 )
-from chakra.et_def.et_def_pb2 import (
+from ....schema.protobuf.et_def_pb2 import (
     NodeType as ChakraNodeType,
 )
-from chakra.third_party.utils.protolib import encodeMessage as encode_message
+from ..third_party.protolib import encodeMessage as encode_message
 
 NODE_ID = 0
 

diff --git a/src/python/jsonizer/jsonizer.py b/src/python/jsonizer/jsonizer.py
@@ -2,15 +2,16 @@
 
 import argparse
 
-from chakra.et_def.et_def_pb2 import (
+from google.protobuf.json_format import MessageToJson
+
+from ....schema.protobuf.et_def_pb2 import (
     GlobalMetadata,
 )
-from chakra.et_def.et_def_pb2 import (
+from ....schema.protobuf.et_def_pb2 import (
     Node as ChakraNode,
 )
-from chakra.third_party.utils.protolib import decodeMessage as decode_message
-from chakra.third_party.utils.protolib import openFileRd as open_file_rd
-from google.protobuf.json_format import MessageToJson
+from ..third_party.protolib import decodeMessage as decode_message
+from ..third_party.protolib import openFileRd as open_file_rd
 
 
 def main() -> None:

diff --git a/src/python/visualizer/visualizer.py b/src/python/visualizer/visualizer.py
@@ -5,12 +5,13 @@
 
 import graphviz
 import networkx as nx
-from chakra.et_def.et_def_pb2 import (
+
+from ....schema.protobuf.et_def_pb2 import (
     GlobalMetadata,
     Node,
 )
-from chakra.third_party.utils.protolib import decodeMessage as decode_message
-from chakra.third_party.utils.protolib import openFileRd as open_file_rd
+from ..third_party.protolib import decodeMessage as decode_message
+from ..third_party.protolib import openFileRd as open_file_rd
 
 
 def escape_label(label: str) -> str: