Revert "Re-arranging the graph architecture"

This reverts commit f145c8317014f950a0df7de4abe075acd149850a.