diff --git a/demos/run_demo.py b/demos/run_demo.py
index 4e6c3bc..89eea18 100644
--- a/demos/run_demo.py
+++ b/demos/run_demo.py
@@ -8,6 +8,7 @@
 import tensorflow.compat.v1 as tf
 import numpy as np
 from graph_optimizer.utils import load_graph, create_node, create_complex_concat_graph
+from graph_optimizer.utils import create_complex_concat_graph
 from graph_optimizer.runner import OptimizationPipeline
 from graph_optimizer.utils.logger import set_log_level, DEBUG, INFO
 
diff --git a/__init__.py b/graph_optimizer/__init__.py
similarity index 100%
rename from __init__.py
rename to graph_optimizer/__init__.py
diff --git a/core.py b/graph_optimizer/core.py
similarity index 100%
rename from core.py
rename to graph_optimizer/core.py
diff --git a/runner.py b/graph_optimizer/runner.py
similarity index 100%
rename from runner.py
rename to graph_optimizer/runner.py
diff --git a/transforms/__init__.py b/graph_optimizer/transforms/__init__.py
similarity index 100%
rename from transforms/__init__.py
rename to graph_optimizer/transforms/__init__.py
diff --git a/transforms/combine/__init__.py b/graph_optimizer/transforms/combine/__init__.py
similarity index 100%
rename from transforms/combine/__init__.py
rename to graph_optimizer/transforms/combine/__init__.py
diff --git a/transforms/combine/concat_combine.py b/graph_optimizer/transforms/combine/concat_combine.py
similarity index 100%
rename from transforms/combine/concat_combine.py
rename to graph_optimizer/transforms/combine/concat_combine.py
diff --git a/transforms/scalar/__init__.py b/graph_optimizer/transforms/scalar/__init__.py
similarity index 100%
rename from transforms/scalar/__init__.py
rename to graph_optimizer/transforms/scalar/__init__.py
diff --git a/transforms/scalar/algebraic_simplify.py b/graph_optimizer/transforms/scalar/algebraic_simplify.py
similarity index 100%
rename from transforms/scalar/algebraic_simplify.py
rename to graph_optimizer/transforms/scalar/algebraic_simplify.py
diff --git a/transforms/scalar/constant_fold.py b/graph_optimizer/transforms/scalar/constant_fold.py
similarity index 100%
rename from transforms/scalar/constant_fold.py
rename to graph_optimizer/transforms/scalar/constant_fold.py
diff --git a/transforms/scalar/cse.py b/graph_optimizer/transforms/scalar/cse.py
similarity index 100%
rename from transforms/scalar/cse.py
rename to graph_optimizer/transforms/scalar/cse.py
diff --git a/transforms/vectorize/__init__.py b/graph_optimizer/transforms/vectorize/__init__.py
similarity index 100%
rename from transforms/vectorize/__init__.py
rename to graph_optimizer/transforms/vectorize/__init__.py
diff --git a/transforms/vectorize/pack_vectorize.py b/graph_optimizer/transforms/vectorize/pack_vectorize.py
similarity index 100%
rename from transforms/vectorize/pack_vectorize.py
rename to graph_optimizer/transforms/vectorize/pack_vectorize.py
diff --git a/utils/__init__.py b/graph_optimizer/utils/__init__.py
similarity index 100%
rename from utils/__init__.py
rename to graph_optimizer/utils/__init__.py
diff --git a/utils/generators.py b/graph_optimizer/utils/generators.py
similarity index 100%
rename from utils/generators.py
rename to graph_optimizer/utils/generators.py
diff --git a/utils/graph_utils.py b/graph_optimizer/utils/graph_utils.py
similarity index 99%
rename from utils/graph_utils.py
rename to graph_optimizer/utils/graph_utils.py
index 211e549..391890f 100644
--- a/utils/graph_utils.py
+++ b/graph_optimizer/utils/graph_utils.py
@@ -6,6 +6,8 @@
 and reusability.
 """
 
+from functools import lru_cache
+
 import os
 import collections
 import numpy as np
@@ -132,6 +134,7 @@ def make_output_shapes_attr(shapes: List[List[int]]) -> attr_value_pb2.AttrValue
 # =======================
 
 
+@lru_cache(maxsize=None)
 def extract_base_name(input_name: str) -> str:
     """
     Extract base node name from input (strip port and control marker).
diff --git a/utils/logger.py b/graph_optimizer/utils/logger.py
similarity index 100%
rename from utils/logger.py
rename to graph_optimizer/utils/logger.py