Remove mock_ttnn

tests/lowering/eltwise/test_add.py

@@ -1,7 +1,7 @@
 import torch
 import torch_ttnn
 import unittest
-from torch_ttnn import ttnn
+import ttnn
 from torch.fx.passes.dialect.common.cse_pass import CSEPass
 
 

tests/lowering/matmul/test_only_add_matmul.py

@@ -1,7 +1,7 @@
 import torch
 import torch_ttnn
 import unittest
-from torch_ttnn import ttnn
+import ttnn
 
 from torch_ttnn.utils import check_with_pcc
 

tests/lowering/tensor_manipulation/test_expand.py

@@ -1,7 +1,7 @@
 import torch
 import torch_ttnn
 import unittest
-from torch_ttnn import ttnn
+import ttnn
 import tt_lib
 from torch_ttnn.utils import (
     DummyTtnnRowMajorLayout,

tests/lowering/test_if.py

@@ -1,7 +1,7 @@
 import torch
 import torch_ttnn
 import unittest
-from torch_ttnn import ttnn
+import ttnn
 from torch.fx.passes.dialect.common.cse_pass import CSEPass
 
 

tests/lowering/test_more_ops.py

@@ -1,7 +1,7 @@
 import torch
 import torch_ttnn
 import unittest
-from torch_ttnn import ttnn
+import ttnn
 import tt_lib
 
 from torch_ttnn.utils import check_with_pcc

tests/models/bert/test_bert.py

@@ -1,7 +1,7 @@
 import torch
 import torch_ttnn
 import unittest
-from torch_ttnn import ttnn
+import ttnn
 import collections
 
 # Load model directly

tests/models/resnet/test_resnet.py

@@ -2,7 +2,7 @@
 import torchvision
 import torch_ttnn
 import unittest
-from torch_ttnn import ttnn
+import ttnn
 import collections
 
 

tests/test_fall_back.py

@@ -1,7 +1,7 @@
 import torch
 import torch_ttnn
 import unittest
-from torch_ttnn import ttnn
+import ttnn
 
 from torch_ttnn.utils import check_with_pcc
 

tests/tools/test_stats.py

@@ -1,7 +1,7 @@
 import os
 import shutil
 import torch
-from torch_ttnn import torch_stat
+import torch_stat
 import unittest
 import json
 from torch.fx.passes.dialect.common.cse_pass import CSEPass

torch_ttnn/__init__.py

@@ -1,8 +1,10 @@
-from .backend import backend
-from .backend import TorchTtnnOption
+from torch_ttnn.backend import backend
+from torch_ttnn.backend import TorchTtnnOption
 
 try:
     import ttnn
-except ImportError:
-    print("ttnn is not installed, use mock_ttnn instead")
-    from . import mock_ttnn as ttnn
+except ImportError as e:
+    print(
+        "ttnn is not installed. Run `python3 -m pip install -r requirements.txt` or `python3 -m pip install -r requirements-dev.txt` if you are developing the compiler"
+    )
+    raise e

torch_ttnn/backend.py

@@ -3,16 +3,11 @@
 from typing import List
 import torch._dynamo
 from functorch.compile import make_boxed_func
+import ttnn
 
 torch._dynamo.config.suppress_errors = False
 torch._dynamo.config.verbose = True
 
-try:
-    import ttnn
-except ImportError:
-    print("ttnn is not installed, use mock_ttnn instead")
-    from . import mock_ttnn as ttnn
-
 
 # The backend for torch.compile that converts a graph to use ttnn.
 # The "option" parameter is a dict that contains one key "torch_ttnn_option".

torch_ttnn/fx_graphviz.py

@@ -4,11 +4,7 @@
 import math
 from collections import defaultdict
 
-try:
-    import ttnn
-except ImportError:
-    print("ttnn is not installed, use mock_ttnn instead")
-    from torch_tnn import mock_ttnn as ttnn
+import ttnn
 
 
 def _tensor_weight(t):

torch_ttnn/passes/graphviz_pass.py

@@ -1,11 +1,7 @@
 import torch
 from torch_ttnn import fx_graphviz
 
-try:
-    import ttnn
-except ImportError:
-    print("ttnn is not installed, use mock_ttnn instead")
-    from torch_ttnn import mock_ttnn as ttnn
+import ttnn
 
 from torch.fx.passes.infra.pass_base import PassBase, PassResult
 

torch_ttnn/passes/lowering/add_data_move_pass.py

@@ -1,16 +1,12 @@
 import torch
+import ttnn
 from torch_ttnn.utils import (
     DummyTtnnUint32,
     DummyTtnnRowMajorLayout,
     DummyTtnnTileLayout,
     DummyDevice,
 )
 
-try:
-    import ttnn
-except ImportError:
-    print("ttnn is not installed, use mock_ttnn instead")
-    from torch_ttnn import mock_ttnn as ttnn
 
 from torch.fx.passes.infra.pass_base import PassBase, PassResult
 

torch_ttnn/passes/lowering/eliminate_data_move_pass.py

@@ -1,10 +1,5 @@
 import torch
-
-try:
-    import ttnn
-except ImportError:
-    print("ttnn is not installed, use mock_ttnn instead")
-    from torch_ttnn import mock_ttnn as ttnn
+import ttnn
 
 from torch.fx.passes.infra.pass_base import PassBase, PassResult
 

torch_ttnn/passes/lowering/permute_reshape_tuple.py

@@ -1,10 +1,5 @@
 import torch
-
-try:
-    import ttnn
-except ImportError:
-    print("ttnn is not installed, use mock_ttnn instead")
-    from torch_ttnn import mock_ttnn as ttnn
+import ttnn
 
 from torch.fx.passes.infra.pass_base import PassBase, PassResult
 

torch_ttnn/passes/lowering/to_tt_pass.py

@@ -1,4 +1,5 @@
 import torch
+import ttnn
 from torch_ttnn.utils import (
     GraphCleanup,
     DummyTtlTensorTensorMemoryLayoutInterleaved,
@@ -9,12 +10,6 @@
 )
 import numpy as np
 
-try:
-    import ttnn
-except ImportError:
-    print("ttnn is not installed, use mock_ttnn instead")
-    from torch_ttnn import mock_ttnn as ttnn
-
 from torch.fx.passes.infra.pass_base import PassBase, PassResult
 import torch.fx.traceback as fx_traceback
 

torch_ttnn/patterns/add.py

@@ -2,12 +2,7 @@
 # This file may not be needed anymore.
 
 import torch
-
-try:
-    import ttnn
-except ImportError:
-    print("ttnn is not installed, use mock_ttnn instead")
-    from torch_ttnn import mock_ttnn as ttnn
+import ttnn
 
 
 # NOTE(yoco) The name `add` must be the same as the name of the function.

torch_ttnn/patterns/mm.py

@@ -2,12 +2,7 @@
 # This file may not be needed anymore.
 
 import torch
-
-try:
-    import ttnn
-except ImportError:
-    print("ttnn is not installed, use mock_ttnn instead")
-    from torch_ttnn import mock_ttnn as ttnn
+import ttnn
 
 
 # NOTE(yoco) The name `matmul` must be the same as the name of the function.