Repository: ise-uiuc/nnsmith (language: python)

Issue:
failed to set new cublas math mode: CUBLAS_STATUS_INVALID_VALUE
1. Switch to `tf` branch.
2. Install tensorflow following [this](https://www.tensorflow.org/install/pip).
3. `pip installl icecream dill`
4. `export TF_FORCE_GPU_ALLOW_GROWTH=true`
5. `python nnsmith/graph_gen.py --max_nodes 5 --viz --framework tensorflow --seed 384525877`

Then we can see this error:

```shell
2022-08-25 13:42:21.370475: E tensorflow/stream_executor/cuda/cuda_blas.cc:197] failed to set new cublas math mode: CUBLAS_STATUS_INVALID_VALUE
2022-08-25 13:42:21.370486: W tensorflow/core/framework/op_kernel.cc:1745] OP_REQUIRES failed at matmul_op_impl.h:438 : INTERNAL: Failed initializing math mode
Traceback (most recent call last):
  File "/home/colin/code/nnsmith/nnsmith/graph_gen.py", line 1326, in <module>
    out_eager = model.run_eagerly(inputs)
  File "/home/colin/code/nnsmith/nnsmith/materialize/tensorflow/__init__.py", line 133, in run_eagerly
    return self.net(**inputs)
  File "/home/colin/miniconda3/envs/py39/lib/python3.9/site-packages/tensorflow/python/util/traceback_utils.py", line 153, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/home/colin/code/nnsmith/nnsmith/materialize/tensorflow/tfnet.py", line 91, in __call__
    out_tensors = instr.fwd_fn(
  File "/home/colin/miniconda3/envs/py39/lib/python3.9/site-packages/keras/utils/traceback_utils.py", line 67, in error_handler
    raise e.with_traceback(filtered_tb) from None
tensorflow.python.framework.errors_impl.InternalError: Exception encountered when calling layer "dense" (type Dense).

Failed initializing math mode [Op:MatMul]

Call arguments received by layer "dense" (type Dense):
  • inputs=tf.Tensor(shape=(1, 1, 1, 1), dtype=float32)
```

Discussions about similar issues on the Internet:
1. https://github.com/tensorflow/tensorflow/issues/9489#issuecomment-297752630
2. Many people think this is caused by an OOM error actually, because by default TF will use nearly all of the memory, which can lead to cublas initialization error. [Way to avoid using most of the memory](https://stackoverflow.com/questions/34199233/how-to-prevent-tensorflow-from-allocating-the-totality-of-a-gpu-memory) The most simple way to disable this behavior of TF is `export TF_FORCE_GPU_ALLOW_GROWTH=1` . But this will not resolve this issue.

Using the seed above `384525877` , we can get this graph:

![image](https://user-images.githubusercontent.com/17642707/186747530-62d8f9c9-59be-4f7f-b9fa-480afb9e8957.png)

Manually building this graph using the program below will not lead to the same error.

```python
from typing import List, Tuple, Dict, Callable
import os
import tensorflow as tf
from tensorflow import keras
# import tensorflow.keras as keras
from keras import layers
from tqdm import tqdm
from icecream import ic

def randn_from_specs(specs: Dict[str, tf.TensorSpec]) -> Dict[str, tf.Tensor]:
    return {
        name: tf.cast(tf.random.normal(shape=spec.shape), dtype=spec.dtype)
        for name, spec in specs.items()
    }

def get_dense(units: int) -> layers.Dense:
    return layers.Dense(units=units, dtype=tf.float32, autocast=False)

class MyModule(tf.Module):
    def __init__(self) -> None:
        self.mlist: List[Callable] = []
        # self._7_linear = layers.Dense(units=2, dtype=tf.float32, autocast=False)
        # self._3_linear = layers.Dense(units=1, dtype=tf.float32, autocast=False)
        # self._5_linear = layers.Dense(units=1, dtype=tf.float32, autocast=False)
        print('init 0...')
        self.mlist.append(get_dense(2))
        print('init 1...')
        self.mlist.append(get_dense(1))
        print('init 2...')
        self.mlist.append(get_dense(1))

    @tf.function # (jit_compile=False)
    def __call__(self, *args, **kwargs) -> Dict[str, tf.Tensor]:
        key2tensor: Dict[int, tf.Tensor] = {}
        key2tensor[0] = kwargs['i0']
        key2tensor[1] = self.mlist[0](key2tensor[0])
        key2tensor[2] = self.mlist[1](key2tensor[1])
        key2tensor[3] = tf.add(key2tensor[1], key2tensor[2])
        key2tensor[4] = self.mlist[2](key2tensor[3])
        key2tensor[5] = tf.add(key2tensor[2], key2tensor[2])
        return {
            "o0": key2tensor[5],
            "o1": key2tensor[4],
        }

if __name__ == '__main__':
    net = MyModule()
    input_dict = randn_from_specs({
        "i0": tf.TensorSpec(shape=[1, 1, 1, 1], dtype=tf.float32),
    })
    ic(input_dict)
    ic(net(**input_dict))
    
```


failed to set new cublas math mode: CUBLAS_STATUS_INVALID_VALUE
1. Switch to `tf` branch.
2. Install tensorflow following [this](https://www.tensorflow.org/install/pip).
3. `pip installl icecream dill`
4. `export TF_FORCE_GPU_ALLOW_GROWTH=true`
5. `python nnsmith/graph_gen.py --max_nodes 5 --viz --framework tensorflow --seed 384525877`

Then we can see this error:

```shell
2022-08-25 13:42:21.370475: E tensorflow/stream_executor/cuda/cuda_blas.cc:197] failed to set new cublas math mode: CUBLAS_STATUS_INVALID_VALUE
2022-08-25 13:42:21.370486: W tensorflow/core/framework/op_kernel.cc:1745] OP_REQUIRES failed at matmul_op_impl.h:438 : INTERNAL: Failed initializing math mode
Traceback (most recent call last):
  File "/home/colin/code/nnsmith/nnsmith/graph_gen.py", line 1326, in <module>
    out_eager = model.run_eagerly(inputs)
  File "/home/colin/code/nnsmith/nnsmith/materialize/tensorflow/__init__.py", line 133, in run_eagerly
    return self.net(**inputs)
  File "/home/colin/miniconda3/envs/py39/lib/python3.9/site-packages/tensorflow/python/util/traceback_utils.py", line 153, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/home/colin/code/nnsmith/nnsmith/materialize/tensorflow/tfnet.py", line 91, in __call__
    out_tensors = instr.fwd_fn(
  File "/home/colin/miniconda3/envs/py39/lib/python3.9/site-packages/keras/utils/traceback_utils.py", line 67, in error_handler
    raise e.with_traceback(filtered_tb) from None
tensorflow.python.framework.errors_impl.InternalError: Exception encountered when calling layer "dense" (type Dense).

Failed initializing math mode [Op:MatMul]

Call arguments received by layer "dense" (type Dense):
  • inputs=tf.Tensor(shape=(1, 1, 1, 1), dtype=float32)
```

Discussions about similar issues on the Internet:
1. https://github.com/tensorflow/tensorflow/issues/9489#issuecomment-297752630
2. Many people think this is caused by an OOM error actually, because by default TF will use nearly all of the memory, which can lead to cublas initialization error. [Way to avoid using most of the memory](https://stackoverflow.com/questions/34199233/how-to-prevent-tensorflow-from-allocating-the-totality-of-a-gpu-memory) The most simple way to disable this behavior of TF is `export TF_FORCE_GPU_ALLOW_GROWTH=1` . But this will not resolve this issue.

Using the seed above `384525877` , we can get this graph:

![image](https://user-images.githubusercontent.com/17642707/186747530-62d8f9c9-59be-4f7f-b9fa-480afb9e8957.png)

Manually building this graph using the program below will not lead to the same error.

```python
from typing import List, Tuple, Dict, Callable
import os
import tensorflow as tf
from tensorflow import keras
# import tensorflow.keras as keras
from keras import layers
from tqdm import tqdm
from icecream import ic

def randn_from_specs(specs: Dict[str, tf.TensorSpec]) -> Dict[str, tf.Tensor]:
    return {
        name: tf.cast(tf.random.normal(shape=spec.shape), dtype=spec.dtype)
        for name, spec in specs.items()
    }

def get_dense(units: int) -> layers.Dense:
    return layers.Dense(units=units, dtype=tf.float32, autocast=False)

class MyModule(tf.Module):
    def __init__(self) -> None:
        self.mlist: List[Callable] = []
        # self._7_linear = layers.Dense(units=2, dtype=tf.float32, autocast=False)
        # self._3_linear = layers.Dense(units=1, dtype=tf.float32, autocast=False)
        # self._5_linear = layers.Dense(units=1, dtype=tf.float32, autocast=False)
        print('init 0...')
        self.mlist.append(get_dense(2))
        print('init 1...')
        self.mlist.append(get_dense(1))
        print('init 2...')
        self.mlist.append(get_dense(1))

    @tf.function # (jit_compile=False)
    def __call__(self, *args, **kwargs) -> Dict[str, tf.Tensor]:
        key2tensor: Dict[int, tf.Tensor] = {}
        key2tensor[0] = kwargs['i0']
        key2tensor[1] = self.mlist[0](key2tensor[0])
        key2tensor[2] = self.mlist[1](key2tensor[1])
        key2tensor[3] = tf.add(key2tensor[1], key2tensor[2])
        key2tensor[4] = self.mlist[2](key2tensor[3])
        key2tensor[5] = tf.add(key2tensor[2], key2tensor[2])
        return {
            "o0": key2tensor[5],
            "o1": key2tensor[4],
        }

if __name__ == '__main__':
    net = MyModule()
    input_dict = randn_from_specs({
        "i0": tf.TensorSpec(shape=[1, 1, 1, 1], dtype=tf.float32),
    })
    ic(input_dict)
    ic(net(**input_dict))
    
```




Discussion / hints:

Fix patch:
diff --git a/.github/actions/setup/action.yaml b/.github/actions/setup/action.yaml
new file mode 100644
index 00000000..4debd10d
--- /dev/null
+++ b/.github/actions/setup/action.yaml
@@ -0,0 +1,44 @@
+name: Setup on Linux
+description: Setup basic conda and Python environment on Linux
+
+inputs:
+  mode:
+    description: whether to install all dependencies
+    required: true
+    default: all
+
+runs:
+  using: composite
+  steps:
+    - uses: actions/checkout@v3
+    - name: install graphviz from apt
+      run: sudo apt install -y graphviz
+      shell: bash
+    - name: set path to dependencies
+      id: set-dep-path
+      run: |
+        if [ ${{ inputs.mode == 'all' }} ]; then
+          echo "::set-output name=dep-path::requirements/**/*.txt"
+        elif [ ${{ inputs.mode == 'none' }} ]; then
+          echo "::set-output name=dep-path::''"
+        else
+          echo "Wrong mode ${{ inputs.mode }}"
+          exit 1
+        fi
+      shell: bash
+    - name: set up Python with cache
+      uses: actions/setup-python@v4
+      id: py
+      with:
+        python-version: '3.10'
+        cache: 'pip'
+        cache-dependency-path: ${{ steps.set-dep-path.outputs.dep-path }}
+    - run: echo "cache hit ${{ steps.py.outputs.cache-hit }}"
+      shell: bash
+    - name: install dependencies
+      run: |
+        if [ ${{ inputs.mode == 'all' }} ]; then
+          bash requirements/install_all.sh
+        fi
+        echo "PYTHONPATH=$PYTHONPATH:${GITHUB_WORKSPACE}" >> $GITHUB_ENV
+      shell: bash
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
new file mode 100644
index 00000000..55dea49e
--- /dev/null
+++ b/.github/workflows/ci.yaml
@@ -0,0 +1,27 @@
+name: CI
+
+on: [push]
+
+jobs:
+  pytest:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - uses: ./.github/actions/setup
+        with:
+          mode: all
+
+      - name: Test core
+        run: |
+          pytest tests/core
+      - name: Test torch
+        run: |
+          pytest tests/torch
+      - name: Test TensorFlow
+        run: |
+          pytest tests/tensorflow
+      - name: Test CLI
+        run: |
+          yes | python nnsmith/cli/model_gen.py model.max_nodes=5 debug.viz=true model.type=torch
+          yes | python nnsmith/cli/model_gen.py model.max_nodes=5 debug.viz=true model.type=onnx
+          yes | python nnsmith/cli/model_gen.py model.max_nodes=5 debug.viz=true model.type=tensorflow
diff --git a/.github/workflows/pre-commit.yaml b/.github/workflows/pre-commit.yaml
new file mode 100644
index 00000000..ec4d9b3a
--- /dev/null
+++ b/.github/workflows/pre-commit.yaml
@@ -0,0 +1,14 @@
+name: pre-commit
+
+on:
+  pull_request:
+  push:
+    branches: [main]
+
+jobs:
+  pre-commit:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - uses: actions/setup-python@v3
+      - uses: pre-commit/action@v3.0.0
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 9f7d40cc..a0c15d88 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -1,11 +1,19 @@
 repos:
+-   repo: https://github.com/pycqa/isort
+    rev: 5.10.1
+    hooks:
+    -   id: isort
+        name: isort (python)
+        args: ["--profile", "black"]
+        exclude: |
+            (?x)^(
+                experiments/.*.py |
+                nnsmith/fuzz.py   |
+                nnsmith/backend_executor.py
+            )$
 -   repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v4.3.0
     hooks:
     -   id: check-yaml
     -   id: end-of-file-fixer
     -   id: trailing-whitespace
--   repo: https://github.com/psf/black
-    rev: 22.6.0
-    hooks:
-    -   id: black
diff --git a/README.md b/README.md
index 59b614bd..e09a7654 100644
--- a/README.md
+++ b/README.md
@@ -36,7 +36,9 @@ We want **code simplicity**: keeping minimal dependencies and focusing on a smal
 Run tests before commit:
 
 ```shell
-pytest tests -s
+# env of torch & tf will conflict so split their unit tests.
+pytest tests/torch -s
+pytest tests/tensorflow -s
 ```
 
 ## Commands
diff --git a/nnsmith/abstract/arith.py b/nnsmith/abstract/arith.py
index 815cadc8..bbb8a1ed 100644
--- a/nnsmith/abstract/arith.py
+++ b/nnsmith/abstract/arith.py
@@ -2,10 +2,8 @@
 
 import z3
 
-
 from nnsmith.error import SanityCheck
 
-
 ARITH_MAX_WIDTH: int = 64
 
 
diff --git a/nnsmith/abstract/dtype.py b/nnsmith/abstract/dtype.py
index 788bbba4..3ca6677d 100644
--- a/nnsmith/abstract/dtype.py
+++ b/nnsmith/abstract/dtype.py
@@ -2,6 +2,7 @@
 
 import numpy as np
 
+
 # TODO(@ganler): add float16 support.
 @unique
 class DType(Enum):
@@ -63,8 +64,22 @@ def from_str(s):
             "bool": DType.bool,
         }[s]
 
+    def numpy(self):
+        return {
+            DType.float16: np.float16,
+            DType.float32: np.float32,
+            DType.float64: np.float64,
+            DType.int8: np.int8,
+            DType.int16: np.int16,
+            DType.int32: np.int32,
+            DType.int64: np.int64,
+            DType.complex64: np.complex64,
+            DType.complex128: np.complex128,
+            DType.bool: np.bool_,
+        }[self]
+
     # TODO(@ganler): put "torchization" in a separate file.
-    def torch(self):
+    def torch(self) -> "torch.dtype":
         import torch
 
         return {
@@ -80,21 +95,8 @@ def torch(self):
             DType.bool: torch.bool,
         }[self]
 
-    def numpy(self):
-        return {
-            DType.float16: np.float16,
-            DType.float32: np.float32,
-            DType.float64: np.float64,
-            DType.int8: np.int8,
-            DType.int16: np.int16,
-            DType.int32: np.int32,
-            DType.int64: np.int64,
-            DType.complex64: np.complex64,
-            DType.complex128: np.complex128,
-            DType.bool: np.bool_,
-        }[self]
-
-    def from_torch(self):
+    @staticmethod
+    def from_torch(dtype) -> "DType":
         import torch
 
         return {
@@ -108,8 +110,29 @@ def from_torch(self):
             torch.complex64: DType.complex64,
             torch.complex128: DType.complex128,
             torch.bool: DType.bool,
+        }[dtype]
+
+    def tensorflow(self) -> "tf.Dtype":
+        import tensorflow as tf
+
+        return {
+            DType.float32: tf.float32,
+            DType.float64: tf.float64,
+            DType.int32: tf.int32,
+            DType.int64: tf.int64,
         }[self]
 
+    @staticmethod
+    def from_tensorflow(dtype) -> "DType":
+        import tensorflow as tf
+
+        return {
+            tf.float32: DType.float32,
+            tf.float64: DType.float64,
+            tf.int32: DType.int32,
+            tf.int64: DType.int64,
+        }[dtype]
+
 
 DTYPE_ALL = [DType.float32, DType.float64, DType.int32, DType.int64, DType.bool]
 DTYPE_NON_BOOLS = [dtype for dtype in DTYPE_ALL if dtype != DType.bool]
diff --git a/nnsmith/abstract/extension.py b/nnsmith/abstract/extension.py
index 151820ed..bef619ca 100644
--- a/nnsmith/abstract/extension.py
+++ b/nnsmith/abstract/extension.py
@@ -1,6 +1,6 @@
-from typing import List, Type
-import types
 import functools
+import types
+from typing import List, Type
 
 from nnsmith.abstract.op import AbsOpBase
 from nnsmith.abstract.tensor import AbsTensor
diff --git a/nnsmith/abstract/op.py b/nnsmith/abstract/op.py
index 067708a1..2534e8ba 100644
--- a/nnsmith/abstract/op.py
+++ b/nnsmith/abstract/op.py
@@ -1,24 +1,24 @@
+import fnmatch
+import os
+import random
 from abc import ABC, abstractmethod
 from copy import deepcopy
-import fnmatch
 from functools import reduce
-import os
-from typing import List, Optional, Tuple, Union, Dict, Type
 from inspect import signature
-import random
+from typing import Dict, List, Optional, Tuple, Type, Union
 
 import z3
 
-from nnsmith.error import SanityCheck, ConstraintCheck
+from nnsmith.abstract.arith import *
 from nnsmith.abstract.dtype import (
-    DType,
     DTYPE_ALL,
-    DTYPE_NON_BOOLS,
     DTYPE_FLOATS,
     DTYPE_INTS,
+    DTYPE_NON_BOOLS,
+    DType,
 )
-from nnsmith.abstract.arith import *
 from nnsmith.abstract.tensor import AbsTensor
+from nnsmith.error import ConstraintCheck, SanityCheck
 
 # There are following types of constraints at this point:
 # 1. Shape variables must be greater than 0;
diff --git a/nnsmith/abstract/tensor.py b/nnsmith/abstract/tensor.py
index 54d3341b..031b15b0 100644
--- a/nnsmith/abstract/tensor.py
+++ b/nnsmith/abstract/tensor.py
@@ -3,9 +3,9 @@
 
 import z3
 
-from nnsmith.error import SanityCheck, ConstraintCheck
-from nnsmith.abstract.dtype import DType
 from nnsmith.abstract.arith import *
+from nnsmith.abstract.dtype import DType
+from nnsmith.error import ConstraintCheck, SanityCheck
 
 
 class AbsTensor:
diff --git a/nnsmith/backends/factory.py b/nnsmith/backends/factory.py
index 1a0e1b05..b96051a2 100644
--- a/nnsmith/backends/factory.py
+++ b/nnsmith/backends/factory.py
@@ -1,17 +1,19 @@
-from abc import ABC, abstractmethod
-from typing import Callable, Dict, Optional, Union, List
 import traceback
+from abc import ABC, abstractmethod
+from typing import Any, Callable, Dict, List, Optional, Union
 
 import numpy as np
-from nnsmith.abstract.dtype import DType
 
+from nnsmith.abstract.dtype import DType
 from nnsmith.abstract.tensor import AbsTensor
-from nnsmith.materialize import BugReport, Oracle, Stage, Symptom, TestCase, Model
 from nnsmith.difftest import assert_allclose
+from nnsmith.materialize import BugReport, Model, Oracle, Stage, Symptom, TestCase
+
+BackendCallable = Callable[[Dict[str, np.ndarray]], Dict[str, np.ndarray]]
 
 
 class BackendFactory(ABC):
-    def __init__(self, device="cpu", optmax=True, catch_process_crash=True):
+    def __init__(self, device="cpu", optmax: bool = False, catch_process_crash=True):
         super().__init__()
         self.device = device
         self.optmax = optmax
@@ -25,9 +27,7 @@ def system_name(self) -> str:
         pass
 
     def __str__(self) -> str:
-        return (
-            f"{self.system_name} (opt={'max' if self.optmax else 'min'}-{self.device})"
-        )
+        return f"{self.system_name} ({self.device}  opt: {self.optmax})"
 
     @staticmethod
     def make_random_input(
@@ -45,9 +45,7 @@ def skip_dtypes(cls) -> List[DType]:
         return []
 
     @abstractmethod
-    def make_backend(
-        self, model: Model
-    ) -> Callable[[Dict[str, np.ndarray]], Dict[str, np.ndarray]]:
+    def make_backend(self, model: Model) -> BackendCallable:
         raise NotImplementedError
 
     def verify_testcase(self, testcase: TestCase) -> Optional[BugReport]:
diff --git a/nnsmith/backends/onnxruntime.py b/nnsmith/backends/onnxruntime.py
index b91e818e..82251586 100644
--- a/nnsmith/backends/onnxruntime.py
+++ b/nnsmith/backends/onnxruntime.py
@@ -1,12 +1,13 @@
-from typing import Callable, Dict
+from typing import Callable, Dict, cast
 
-from multipledispatch import dispatch
 import numpy as np
 import onnx
 import onnxruntime as ort
+from multipledispatch import dispatch
 
-from nnsmith.macro import NNSMITH_ORT_INTRA_OP_THREAD
 from nnsmith.backends import BackendFactory
+from nnsmith.backends.factory import BackendCallable
+from nnsmith.macro import NNSMITH_ORT_INTRA_OP_THREAD
 from nnsmith.materialize.onnx import ONNXModel
 
 OPT_LEVELS = [
@@ -38,8 +39,9 @@ def system_name(self) -> str:
 
     @dispatch(ONNXModel)
     def make_backend(
-        self, model: ONNXModel
-    ) -> Callable[[Dict[str, np.ndarray]], Dict[str, np.ndarray]]:
+        self,
+        model: ONNXModel,
+    ) -> BackendCallable:
         sess_options = ort.SessionOptions()
         sess_options.graph_optimization_level = self.opt_level
         # https://github.com/microsoft/onnxruntime/issues/8313
diff --git a/nnsmith/backends/tensorrt.py b/nnsmith/backends/tensorrt.py
index c785ed29..d4ee0c33 100644
--- a/nnsmith/backends/tensorrt.py
+++ b/nnsmith/backends/tensorrt.py
@@ -1,19 +1,19 @@
-from dataclasses import dataclass
 import warnings
+from dataclasses import dataclass
 from typing import List
 
-import onnx
 import numpy as np
-import tensorrt as trt
+import onnx
 import pycuda.driver as cuda
-from pycuda.driver import DeviceAllocation
+import tensorrt as trt
 from multipledispatch import dispatch
+from pycuda.driver import DeviceAllocation
 
-from nnsmith.abstract.op import AbsOpBase
-from nnsmith.abstract.tensor import AbsTensor
-from nnsmith.abstract.dtype import DType
 from nnsmith.abstract.arith import *
+from nnsmith.abstract.dtype import DType
 from nnsmith.abstract.extension import patch_requires
+from nnsmith.abstract.op import AbsOpBase
+from nnsmith.abstract.tensor import AbsTensor
 from nnsmith.backends import BackendFactory
 from nnsmith.materialize.onnx import ONNXModel
 
diff --git a/nnsmith/backends/tflite.py b/nnsmith/backends/tflite.py
new file mode 100644
index 00000000..b2ae96f3
--- /dev/null
+++ b/nnsmith/backends/tflite.py
@@ -0,0 +1,162 @@
+from os import PathLike
+from typing import Callable, Dict
+
+import numpy as np
+import tensorflow as tf  # type: ignore
+from multipledispatch import dispatch
+
+from nnsmith.backends.factory import BackendCallable, BackendFactory
+from nnsmith.materialize.tensorflow import (
+    TFModel,
+    TFNetCallable,
+    TFNetInputDict,
+    TFNetOutDict,
+    np_dict_from_tf,
+    tf_dict_from_np,
+)
+
+
+class TFLiteRunner:
+    def __init__(self, tfnet_callable: TFNetCallable) -> None:
+        self.tfnet_callable = tfnet_callable
+
+    def __call__(self, input: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
+        return self.tfnet_callable(**input)
+        # https://github.com/tensorflow/tensorflow/issues/34536#issuecomment-565632906
+        # TFLite doesn't support NVIDIA GPU.
+        # It can automatically convert input args to np.ndarray, and it outputs np.ndarray.
+        tf_input = tf_dict_from_np(input)
+        tf_output = self.tfnet_callable(**tf_input)
+        return np_dict_from_tf(tf_output)
+
+
+class TFLiteFactory(BackendFactory):
+    """Factory to build TFLite backend.
+    Convertion Graph:
+        TFModel & one concrete function
+                    v  TFLiteConverter with configs
+        TFLite model content in bytes -- (f.write) --> model.tflite file
+                    v  tf.lite.Interpreter  <-- (f.read) <---
+        TFLite Python Callable
+    """
+
+    def __init__(self, device, optmax, **kwargs) -> None:
+        super().__init__(device, optmax, **kwargs)
+
+    @property
+    def system_name(self) -> str:
+        "tflite"
+
+    @dispatch(TFModel)
+    def make_backend(
+        self,
+        model: TFModel,
+    ) -> BackendCallable:
+        """Create TFLite callable from a concrete function.
+        TFModel is required because functions have a weak reference to Variables, which are stored in `tf.Module`.
+        Ref: https://www.tensorflow.org/api_docs/python/tf/lite/TFLiteConverter#from_concrete_functions
+        """
+        return self.make_backend_from_content(self.make_content(model))
+
+    def make_backend_from_path(
+        self,
+        path: PathLike,
+    ) -> BackendCallable:
+        """Create TFLite callable from path of a TF SavedModel.
+        Ref: https://www.tensorflow.org/api_docs/python/tf/lite/TFLiteConverter#from_saved_model
+        """
+        return self.make_backend_from_content(
+            self.make_content(path),
+        )
+
+    @dispatch(bytes)
+    def make_backend_from_content(self, content: bytes) -> BackendCallable:
+        # Ref: https://www.tensorflow.org/api_docs/python/tf/lite/Interpreter
+        interpreter = tf.lite.Interpreter(model_content=content)
+        return TFLiteRunner(interpreter.get_signature_runner())
+
+    @dispatch(TFModel)
+    def make_content(
+        self,
+        model: TFModel,
+    ) -> bytes:
+        """Create TFLite content from a concrete function.
+        TFModel is required because functions have a weak reference to Variables, which are stored in `tf.Module`.
+        Ref: https://www.tensorflow.org/api_docs/python/tf/lite/TFLiteConverter#from_concrete_functions
+        """
+        converter = tf.lite.TFLiteConverter.from_concrete_functions(
+            funcs=[model.concrete_net()],
+            trackable_obj=model.net,
+        )
+        return self._tflite_content_from_converter(converter)
+
+    def make_content_with_func(
+        self,
+        model: TFModel,
+        concrete_func: TFNetCallable,
+    ) -> bytes:
+        """Create TFLite content from a concrete function.
+        TFModel is required because functions have a weak reference to Variables, which are stored in `tf.Module`.
+        Ref: https://www.tensorflow.org/api_docs/python/tf/lite/TFLiteConverter#from_concrete_functions
+        """
+        converter = tf.lite.TFLiteConverter.from_concrete_functions(
+            funcs=[concrete_func],
+            trackable_obj=model.net,
+        )
+        return self._tflite_content_from_converter(converter)
+
+    def make_backend_with_func(
+        self,
+        model: TFModel,
+        concrete_func: TFNetCallable,
+    ) -> BackendCallable:
+        """Create TFLite callable from a concrete function.
+        TFModel is required because functions have a weak reference to Variables, which are stored in `tf.Module`.
+        Ref: https://www.tensorflow.org/api_docs/python/tf/lite/TFLiteConverter#from_concrete_functions
+        """
+        return self.make_backend_from_content(
+            self.make_content_with_func(model, concrete_func)
+        )
+
+    @dispatch(str)
+    def make_content(
+        self,
+        path: PathLike,
+    ) -> bytes:
+        """Create TFLite content from path of a saved model.
+        Ref: https://www.tensorflow.org/api_docs/python/tf/lite/TFLiteConverter#from_saved_model
+        """
+        converter = tf.lite.TFLiteConverter.from_saved_model(path)
+        return self._tflite_content_from_converter(converter)
+
+    def _tflite_content_from_converter(
+        self,
+        converter: tf.lite.TFLiteConverter,
+    ) -> bytes:
+        """Configure TFLite converter and create callable from it.
+        converter configuarations: https://www.tensorflow.org/api_docs/python/tf/lite/TFLiteConverter#attributes_1
+        """
+        # Ref: https://www.tensorflow.org/api_docs/python/tf/lite/TargetSpec
+        converter.target_spec.supported_ops = [
+            tf.lite.OpsSet.TFLITE_BUILTINS,  # enable TensorFlow Lite ops.
+            tf.lite.OpsSet.SELECT_TF_OPS,  # enable TensorFlow ops.
+        ]
+        # Ref: https://www.tensorflow.org/api_docs/python/tf/lite/Optimize
+        converter.optimizations = {
+            tf.lite.Optimize.DEFAULT,
+            tf.lite.Optimize.EXPERIMENTAL_SPARSITY,
+        }
+        # converter.allow_custom_ops = True
+        tflite_bytes = converter.convert()
+        return tflite_bytes
+
+    def dump_backend(self, path: PathLike, content: bytes) -> None:
+        with open(path, "wb") as f:
+            f.write(content)
+
+    def load_content(self, path: PathLike) -> bytes:
+        with open(path, "rb") as f:
+            return f.read()
+
+    def load_backend(self, path: PathLike) -> BackendCallable:
+        return self.make_backend_from_content(self.load_content(path))
diff --git a/nnsmith/backends/tvm.py b/nnsmith/backends/tvm.py
index dc8468a9..149f900e 100644
--- a/nnsmith/backends/tvm.py
+++ b/nnsmith/backends/tvm.py
@@ -1,6 +1,6 @@
 import tvm
-from tvm import relay
 from multipledispatch import dispatch
+from tvm import relay
 
 from nnsmith.backends import BackendFactory
 from nnsmith.materialize.onnx import ONNXModel
diff --git a/nnsmith/cli/model_gen.py b/nnsmith/cli/model_gen.py
index c44a1aca..b52bcb16 100644
--- a/nnsmith/cli/model_gen.py
+++ b/nnsmith/cli/model_gen.py
@@ -4,7 +4,8 @@
 import hydra
 from omegaconf import DictConfig
 
-from nnsmith.graph_gen import random_model_gen, viz, concretize_graph, make_schedule
+from nnsmith.graph_gen import concretize_graph, random_model_gen, viz
+from nnsmith.materialize import Schedule
 
 
 @hydra.main(version_base=None, config_path="../config", config_name="main")
@@ -16,8 +17,7 @@ def main(cfg: DictConfig):
     print(f"Using seed {seed}")
 
     # TODO(@ganler): skip operators outside of model gen with `cfg[exclude]`
-    from nnsmith.materialize import TestCase
-    from nnsmith.materialize import Model
+    from nnsmith.materialize import Model, TestCase
     from nnsmith.util import mkdir
 
     model_cfg = cfg["model"]
@@ -43,7 +43,7 @@ def main(cfg: DictConfig):
         gen.abstract_graph, gen.tensor_dataflow, gen.get_solutions()
     )
 
-    schedule = make_schedule(fixed_graph, concrete_abstensors)
+    schedule = Schedule.init(fixed_graph, concrete_abstensors)
 
     model = ModelType.from_schedule(schedule)
     model.refine_weights()  # either random generated or gradient-based.
diff --git a/nnsmith/graph_gen.py b/nnsmith/graph_gen.py
index 32b6395c..dd1ae22a 100644
--- a/nnsmith/graph_gen.py
+++ b/nnsmith/graph_gen.py
@@ -1,24 +1,23 @@
-from collections import defaultdict, namedtuple
+import copy
 import math
-import textwrap
-from typing import Dict, Tuple, List, Set, Type
-from inspect import signature
-import traceback
+import os
 import random
+import textwrap
 import time
-import os
-import copy
-import warnings
+import traceback
 import uuid
+import warnings
+from collections import defaultdict, namedtuple
+from inspect import signature
+from typing import Dict, List, Set, Tuple, Type
 
 import networkx as nx
 import z3
 
-from nnsmith.util import set_seed
-from nnsmith.error import SanityCheck, ConstraintError
-from nnsmith.abstract.op import *
 from nnsmith.abstract.op import __MAX_RANK__ as __MAX_RANK__
-
+from nnsmith.abstract.op import *
+from nnsmith.error import ConstraintError, SanityCheck
+from nnsmith.util import set_seed
 
 NNSMITH_LIMNF_V = os.getenv("NNSMITH_LIMNF_V", "0")
 assert NNSMITH_LIMNF_V in ["0", "1"]
@@ -35,57 +34,12 @@ class RequiredDimNotFound(Exception):
 
 TensorCtx = namedtuple("TensorCtx", ["op_id", "type", "output_idx"])
 # Every tensor is represented by a (unique) integer key.
-Instruction: Tuple[AbsOpBase, List[int], List[int]] = namedtuple(
-    "Instruction", ["op", "inputs", "outputs"]
-)
-
-
-# Minimal information for constructing a graph.
-class Schedule:
-    def __init__(self, instructions, input_keys, leaf_keys, key2type):
-        self.instructions: List[Instruction] = instructions
-        self.input_keys: List[int] = input_keys
-        self.leaf_keys: List[int] = leaf_keys
-        self.key2type: Dict[int, AbsTensor] = key2type
-
-
-def make_schedule(graph: nx.MultiDiGraph, key2type: Dict[int, AbsTensor]):
-    # The input graph should be a concretized graph.
-    instructions: List[Instruction] = []
-    input_keys = []
-    user_keys = set()
-
-    # freeze node with static attributes in label;
-    for node_id in nx.topological_sort(graph):
-        node = graph.nodes[node_id]
-        op = node["op"]
-
-        if isinstance(op, Input):
-            input_keys.append(node["otensor_idx"][0])
-
-        for used_idx in node["itensor_idx"]:
-            user_keys.add(used_idx)
-
-        # TODO(@ganler): Better name than "otensor_idx"
-        # TODO(@ganler): Add refcnt or last ref mechanism to save memory
-        instructions.append(
-            Instruction(
-                op=op,
-                inputs=node["itensor_idx"],
-                outputs=node["otensor_idx"],
-            )
-        )
-
-    # simplify the statements above
-    leaf_keys = [key for key in key2type if key not in user_keys]
-
-    return Schedule(instructions, input_keys, leaf_keys, key2type)
 
 
 def concretize_graph(
     graph: nx.MultiDiGraph, dataflow: List[TensorCtx], model: z3.ModelRef
-) -> Tuple[nx.MultiDiGraph,]:
-    concrete_shapes = {}
+) -> Tuple[nx.MultiDiGraph, Dict[int, AbsTensor]]:
+    concrete_shapes: Dict[int, AbsTensor] = {}
 
     # freeze node with static attributes in label;
     for node_id in nx.topological_sort(graph):
@@ -1137,4 +1091,4 @@ def viz(G, filename: str = None):
     elif filename.endswith("svg"):
         nx.drawing.nx_pydot.to_pydot(G).write_svg(filename)
     else:
-        raise ValueError(f"Unsupported image format: {fmt}")
+        raise ValueError(f"Unsupported image format: {filename}")
diff --git a/nnsmith/materialize/__init__.py b/nnsmith/materialize/__init__.py
index b3b296f6..697b87e0 100644
--- a/nnsmith/materialize/__init__.py
+++ b/nnsmith/materialize/__init__.py
@@ -1,19 +1,113 @@
-import pickle
-import os
 import json
+import os
+import pickle
 from abc import ABC, abstractmethod
+from collections import namedtuple
+from dataclasses import dataclass
 from enum import Enum
-from typing import Dict, List, Type
+from os import PathLike
+from typing import Any, Dict, List, Set, Tuple, Type
 
+import networkx as nx
 import numpy as np
+from multipledispatch import dispatch
 
-from nnsmith.abstract.op import AbsOpBase
+from nnsmith.abstract.op import AbsOpBase, Constant, Input
 from nnsmith.abstract.tensor import AbsTensor
-from nnsmith.graph_gen import Schedule
+from nnsmith.error import SanityCheck
+
+
+def framework_operator_impl(
+    framework_realizable_ops: List[Type[AbsOpBase]],
+    all_framework_ops: List[Type[AbsOpBase]],
+    op_type: AbsOpBase,
+    *args,
+    **kwargs,
+):
+    """When implementing `forward_fn` for an operator class, add this operator into all_framework_ops list.
+
+    Usage:
+        In `forward.py`, define `operator_impl = partial(framework_operator_impl, FW_REALIZABLE_OPS, ALL_FM_OPS)`.
+        Then add `@operator_impl(OpClass)` when implementing `forward_fn` for `OpClass`.
+
+    Args:
+        framework_realizable_ops (List[Type[AbsOpBase]]): all realizable ops in the framework. Usually it can be obtained by FULL_OPERATOR_SETS["core"].union(FULL_OPERATOR_SETS["framework_name"])
+        all_framework_ops (List[Type[AbsOpBase]]): list of operator classes that are implemented `forward_fn` in the framework.
+        op_type (AbsOpBase): operator class
+    """
+    SanityCheck.true(
+        issubclass(op_type, AbsOpBase),
+        f"Decorator operator_impl takes AbsOpBase subclass, but got {op_type}",
+    )
+    if op_type is not Constant:  # Constant comes from placeholder.
+        dispatchables = [
+            rtype for rtype in framework_realizable_ops if issubclass(rtype, op_type)
+        ]
+        for rtype in dispatchables:
+            all_framework_ops.append(rtype)
+
+        SanityCheck.true(
+            len(dispatchables) != 0,
+            f"Decorator operator_impl only take types decorated by `mark_realize`, but got {op_type}",
+        )
+    return dispatch(op_type, *args, **kwargs)
+
+
+Instruction: Tuple[AbsOpBase, List[int], List[int]] = namedtuple(
+    "Instruction", ["op", "inputs", "outputs"]
+)
+
+
+@dataclass
+class Schedule:
+    """Minimal information for constructing a graph."""
+
+    instructions: List[Instruction]
+    input_keys: List[int]
+    leaf_keys: List[int]
+    key2type: Dict[int, AbsTensor]
+
+    @staticmethod
+    def init(graph: nx.MultiDiGraph, key2type: Dict[int, AbsTensor]) -> "Schedule":
+        # The input graph should be a concretized graph.
+        instructions: List[Instruction] = []
+        input_keys = []
+        user_keys = set()
+
+        # freeze node with static attributes in label;
+        for node_id in nx.topological_sort(graph):
+            node = graph.nodes[node_id]
+            op = node["op"]
+
+            if isinstance(op, Input):
+                input_keys.append(node["otensor_idx"][0])
+
+            for used_idx in node["itensor_idx"]:
+                user_keys.add(used_idx)
+
+            # TODO(@ganler): Better name than "otensor_idx"
+            # TODO(@ganler): Add refcnt or last ref mechanism to save memory
+            instructions.append(
+                Instruction(
+                    op=op,
+                    inputs=node["itensor_idx"],
+                    outputs=node["otensor_idx"],
+                )
+            )
+
+        # simplify the statements above
+        leaf_keys = [key for key in key2type if key not in user_keys]
+
+        return Schedule(instructions, input_keys, leaf_keys, key2type)
 
 
 class Oracle:
-    def __init__(self, input, output, provider: str = "unknown") -> None:
+    def __init__(
+        self,
+        input: Dict[str, np.ndarray],
+        output: Dict[str, np.ndarray],
+        provider: str = "unknown",
+    ) -> None:
         self.input: Dict[str, np.ndarray] = input
         self.output: Dict[str, np.ndarray] = output
         self._provider = provider
@@ -29,7 +123,7 @@ def name() -> str:
     def provider(self) -> str:
         return self._provider
 
-    def dump(self, path: str) -> None:
+    def dump(self, path: PathLike) -> None:
         with open(path, "wb") as f:
             to_dump = {
                 "input": self.input,
@@ -39,7 +133,7 @@ def dump(self, path: str) -> None:
             pickle.dump(to_dump, f)
 
     @staticmethod
-    def load(path: str) -> "Oracle":
+    def load(path: PathLike) -> "Oracle":
         with open(path, "rb") as f:
             to_load = pickle.load(f)
             return Oracle(to_load["input"], to_load["output"], to_load["provider"])
@@ -58,21 +152,21 @@ def output_like(self) -> Dict[str, AbsTensor]:
 
     @staticmethod
     @abstractmethod
-    def from_schedule(self, instructions: Schedule, **kwargs) -> "Model":
+    def from_schedule(schedule: Schedule, **kwargs) -> "Model":
         pass
 
     @staticmethod
     @abstractmethod
-    def load(path) -> "Model":
+    def load(path: PathLike) -> "Model":
         pass
 
     @abstractmethod
-    def dump(self, path):
+    def dump(self, path: PathLike) -> None:
         pass
 
     @property
     @abstractmethod
-    def native_model(self):
+    def native_model(self) -> Any:
         pass
 
     @staticmethod
@@ -93,7 +187,12 @@ def refine_weights(self) -> None:
         pass
 
     @staticmethod
-    def name_prefix():
+    @abstractmethod
+    def operators() -> List[Type[AbsOpBase]]:
+        pass
+
+    @staticmethod
+    def name_prefix() -> str:
         return "model"
 
     @staticmethod
@@ -114,6 +213,10 @@ def init(name) -> Type["Model"]:
             from nnsmith.materialize.onnx import ONNXModel
 
             return ONNXModel
+        elif name == "tensorflow":
+            from nnsmith.materialize.tensorflow import TFModel
+
+            return TFModel
         else:
             raise ValueError(f"Unknown or unsupported model type: {name}")
 
@@ -124,7 +227,9 @@ def __init__(self, model: Model, oracle: Oracle):
         self.model = model
 
     @staticmethod
-    def load(model_type: Model, root_folder: str, allow_no_oracle=False) -> "TestCase":
+    def load(
+        model_type: Type[Model], root_folder: str, allow_no_oracle=False
+    ) -> "TestCase":
         model_path = os.path.join(
             root_folder, model_type.name_prefix() + model_type.name_suffix()
         )
diff --git a/nnsmith/materialize/onnx/__init__.py b/nnsmith/materialize/onnx/__init__.py
index 47cd10c1..21e9016a 100644
--- a/nnsmith/materialize/onnx/__init__.py
+++ b/nnsmith/materialize/onnx/__init__.py
@@ -1,22 +1,23 @@
-from typing import Dict, Tuple, Union
-import warnings
-import random
 import os
+import random
+import warnings
 from io import BytesIO
+from os import PathLike
+from typing import Dict, Tuple, Union
 
-import torch
-import torch.onnx
 import onnx
 import onnx.checker
 import onnx.helper
-from onnx.tools import update_model_dims
+import torch
+import torch.onnx
 from onnx.external_data_helper import load_external_data_for_model
+from onnx.tools import update_model_dims
 
 from nnsmith.abstract.dtype import DType
 from nnsmith.abstract.op import AbsTensor
-from nnsmith.graph_gen import Schedule
-from nnsmith.materialize.torch import TorchModel, SymbolNet
 from nnsmith.macro import onnx2external_data_dir
+from nnsmith.materialize import Schedule
+from nnsmith.materialize.torch import SymbolNet, TorchModel
 
 
 def create_deadcode_onnx(onnx_model, name_mask) -> onnx.ModelProto:
@@ -107,7 +108,6 @@ def analyze_onnx_io(
     model: onnx.ModelProto,
 ) -> Tuple[Dict[str, AbsTensor], Dict[str, AbsTensor]]:
     """Return input and output types in dictionaries."""
-    print(model)
     input_types = {}
     output_types = {}
     for input in model.graph.input:
@@ -214,7 +214,7 @@ def input_like(self) -> Dict[str, AbsTensor]:
     def output_like(self) -> Dict[str, AbsTensor]:
         return self.masked_output_like
 
-    def dump(self, path: str) -> None:
+    def dump(self, path: PathLike) -> None:
         if self.with_torch:
             TorchModel.dump(
                 self, path.replace(self.name_suffix(), TorchModel.name_suffix())
@@ -224,7 +224,7 @@ def dump(self, path: str) -> None:
         onnx.save(self.onnx_model, path)
 
     @staticmethod
-    def load(path: str) -> "ONNXModel":
+    def load(path: PathLike) -> "ONNXModel":
         ret = ONNXModel()
         ret.onnx_model = onnx.load(path)
 
diff --git a/nnsmith/materialize/tensorflow/__init__.py b/nnsmith/materialize/tensorflow/__init__.py
new file mode 100644
index 00000000..f1b7be95
--- /dev/null
+++ b/nnsmith/materialize/tensorflow/__init__.py
@@ -0,0 +1,223 @@
+from __future__ import annotations
+
+import os
+import pickle
+from os import PathLike
+from typing import Any, Callable, Dict, List, Tuple, Type
+
+import numpy as np
+import tensorflow as tf  # type: ignore
+from multipledispatch import dispatch  # type: ignore
+
+
+def configure_tensorflow():
+    # https://github.com/tensorflow/tensorflow/issues/57359
+    # tf.config.experimental.enable_tensor_float_32_execution(False)
+    for gpu in tf.config.experimental.list_physical_devices("GPU"):
+        tf.config.experimental.set_memory_growth(gpu, True)
+
+
+configure_tensorflow()
+
+from nnsmith.abstract.op import AbsOpBase, AbsTensor
+from nnsmith.materialize import Model, Oracle, Schedule
+from nnsmith.materialize.tensorflow.forward import ALL_TF_OPS
+from nnsmith.materialize.tensorflow.tfnet import TFNet, TFNetOutDict
+
+TFNetInputDict = Dict[str, tf.Tensor]
+TFNetCallable = Callable[..., TFNetOutDict]
+
+
+@dispatch(dict)
+def randn_from_specs(specs: Dict[str, tf.TensorSpec]) -> Dict[str, tf.Tensor]:
+    return {
+        name: tf.cast(tf.random.normal(shape=spec.shape), dtype=spec.dtype)
+        for name, spec in specs.items()
+    }
+
+
+def np_dict_from_tf(x: Dict[str, tf.Tensor]) -> Dict[str, np.ndarray]:
+    return {key: value.numpy() for key, value in x.items()}
+
+
+def tf_dict_from_np(x: Dict[str, np.ndarray]) -> Dict[str, tf.Tensor]:
+    return {key: tf.convert_to_tensor(value) for key, value in x.items()}
+
+
+def assert_dict_eq_tf(x: Dict[str, tf.Tensor], y: Dict[str, tf.Tensor]) -> None:
+    for key in x:
+        x_v, y_v = x[key], y[key]
+        assert tf.less_equal(
+            tf.reduce_max(tf.abs(x_v - y_v)),
+            tf.cast(1e-3, dtype=x_v.dtype),
+        ), f"Tensors are NOT equal: x[{key}] = {x_v} != {y_v} = y[{key}]"
+
+
+class TFModel(Model):
+    """Wrapper class of TFNet (tf.Module)
+    It only stores net. Other information like I/O info are temporarily generated.
+    Other values like I/O values should be managed by the user.
+    """
+
+    def __init__(
+        self,
+        schedule: Schedule,
+        verbose: bool = False,
+    ) -> None:
+        """Must provide a schedule to avoid NoneType errors"""
+        super().__init__()
+        self.net: TFNet = TFNet(
+            schedule=schedule,
+            verbose=verbose,
+        )
+
+    def set_verbose(self, verbose: bool = True) -> None:
+        self.net.verbose = verbose
+
+    @staticmethod
+    def from_schedule(schedule: Schedule, **kwargs) -> "Model":
+        return TFModel(schedule, kwargs.get("verbose", False))
+
+    @property
+    def native_model(self) -> TFNet:
+        return self.net
+
+    @property
+    def input_like(self) -> Dict[str, AbsTensor]:
+        return {
+            f"i{i_inp}": self.net.schedule.key2type[key]
+            for i_inp, key in enumerate(self.net.schedule.input_keys)
+        }
+
+    @property
+    def output_like(self) -> Dict[str, AbsTensor]:
+        return {
+            f"o{i_out}": self.net.schedule.key2type[key]
+            for i_out, key in enumerate(self.net.schedule.leaf_keys)
+        }
+
+    @property
+    def input_specs(self) -> Dict[str, tf.TensorSpec]:
+        ret: Dict[str, tf.TensorSpec] = {}
+        for i_inp, key in enumerate(self.net.schedule.input_keys):
+            abs_tensor = self.net.schedule.key2type[key]
+            ret[f"i{i_inp}"] = tf.TensorSpec(
+                abs_tensor.shape, abs_tensor.dtype.tensorflow(), f"i{i_inp}"
+            )
+        return ret
+
+    @staticmethod
+    def name_suffix() -> str:
+        return ""
+
+    def make_oracle(
+        self, inputs: Dict[str, tf.Tensor | tf.TensorSpec] = None
+    ) -> Oracle:
+        if inputs is None or isinstance(inputs["i0"], tf.TensorSpec):
+            input_dict = self.random_inputs()
+        else:
+            input_dict = inputs
+        output_dict = self.run_eagerly(input_dict)
+
+        input_dict = np_dict_from_tf(input_dict)
+        output_dict = np_dict_from_tf(output_dict)
+
+        return Oracle(input_dict, output_dict)
+
+    def dump(self, path: PathLike = "saved_tfmodel") -> None:
+        os.makedirs(path, exist_ok=True)
+        # schedule.pkl
+        with open(os.path.join(path, TFModel.schedule_pkl_name()), "wb") as f:
+            pickle.dump(self.net.schedule, f)
+        # tfnet
+        concrete_net = self.concrete_net(self.input_specs)
+        tf.saved_model.save(
+            self.net,
+            os.path.join(path, TFModel.tfnet_dir_name()),
+            signatures=concrete_net,
+        )
+
+    def dump_with_oracle(
+        self,
+        path: PathLike = "saved_tfmodel",
+        inputs: Dict[str, tf.Tensor | tf.TensorSpec] = None,
+    ) -> None:
+        self.dump(path)
+        oracle = self.make_oracle(inputs)
+        oracle.dump(os.path.join(path, Oracle.name()))
+
+    def dump_tfnet(
+        self,
+        path: PathLike = "saved_tfnet",
+        inputs: Dict[str, tf.Tensor | tf.TensorSpec] = None,
+    ) -> TFNetCallable:
+        concrete_net = self.concrete_net(inputs)
+        tf.saved_model.save(self.net, path, signatures=concrete_net)
+        return concrete_net
+
+    @staticmethod
+    def load_tfnet(
+        saved_dir: str = "saved_tfnet",
+    ) -> TFNetCallable:
+        return tf.saved_model.load(saved_dir)
+
+    @staticmethod
+    def load(path: PathLike = "saved_tfmodel") -> "TFModel":
+        with open(os.path.join(path, TFModel.schedule_pkl_name()), "rb") as f:
+            schedule: Schedule = pickle.load(f)
+        model = TFModel(schedule)
+        return model
+
+    @staticmethod
+    def load_with_oracle(
+        path: PathLike = "saved_tfmodel",
+    ) -> Tuple["TFModel", TFNetInputDict, TFNetOutDict]:
+        model = TFModel.load(path)
+        oracle = Oracle.load(os.path.join(path, Oracle.name()))
+        input_dict = {name: tf.convert_to_tensor(v) for name, v in oracle.input.items()}
+        output_dict = {
+            name: tf.convert_to_tensor(v) for name, v in oracle.output.items()
+        }
+        return model, input_dict, output_dict
+
+    @staticmethod
+    def schedule_pkl_name():
+        return "schedule.pkl"
+
+    @staticmethod
+    def in_out_pkl_name():
+        return "in_out.pkl"
+
+    @staticmethod
+    def tfnet_dir_name():
+        return "tfnet"
+
+    def random_inputs(self) -> TFNetInputDict:
+        return {
+            spec.name: tf.cast(
+                tf.random.normal(
+                    shape=spec.shape,
+                    seed=None,
+                ),
+                dtype=spec.dtype,
+            )
+            for spec in self.input_specs.values()
+        }
+
+    def concrete_net(
+        self, inputs: Dict[str, tf.Tensor | tf.TensorSpec] | None = None
+    ) -> TFNetCallable:
+        if inputs is None:
+            inputs = self.input_specs
+        return self.net.__call__.get_concrete_function(**inputs)
+
+    def refine_weights(self) -> None:
+        pass
+
+    def run_eagerly(self, inputs: TFNetInputDict) -> TFNetOutDict:
+        tf.config.run_functions_eagerly(True)
+        return self.net(**inputs)
+
+    @staticmethod
+    def operators() -> List[Type[AbsOpBase]]:
+        return list(ALL_TF_OPS)
diff --git a/nnsmith/materialize/tensorflow/dialect.py b/nnsmith/materialize/tensorflow/dialect.py
new file mode 100644
index 00000000..5c0257fc
--- /dev/null
+++ b/nnsmith/materialize/tensorflow/dialect.py
@@ -0,0 +1,47 @@
+from typing import List, Tuple, Union
+
+from nnsmith.abstract.arith import *
+from nnsmith.abstract.dtype import DType
+from nnsmith.abstract.op import UnaryOpBase, int_from, mark_materialize
+from nnsmith.abstract.tensor import AbsTensor
+from nnsmith.error import ConstraintCheck
+
+
+@mark_materialize("tensorflow")
+class Dense(UnaryOpBase):
+    in_dtypes = [(DType.float32,), (DType.float64,)]
+    out_dtypes = [(DType.float32,), (DType.float64,)]
+
+    def __init__(self, ifeat: Union[int, z3.ExprRef], ofeat: Union[int, z3.ExprRef]):
+        super().__init__()
+        self.ifeat = ifeat
+        self.ofeat = ofeat
+        self.inp_ranks = [
+            int_from(2)
+        ]  # NOTE: tensorflow Dense layer requires an input with batch as its first axis
+        # at least one dim. cannot be zranks_all()
+        self.out_ranks = [int_from(2)]
+
+    def type_transfer(self, input_shapes: List[AbsTensor]) -> List[AbsTensor]:
+        assert len(input_shapes) == 1, "Linear only takes one input, but got {}".format(
+            len(input_shapes)
+        )
+        return [
+            AbsTensor(
+                shape=[*input_shapes[0].shape[:-1], self.ofeat],
+                dtype=input_shapes[0].dtype,
+            )
+        ]
+
+    def requires(self, input_shapes: List[AbsTensor]) -> List[z3.ExprRef]:
+        ConstraintCheck.true(input_shapes[0].ndims >= 2)
+        return [
+            nnsmith_ge(self.ifeat, 1),
+            nnsmith_ge(self.ofeat, 1),
+            nnsmith_eq(input_shapes[0].shape[-1], self.ifeat),
+        ]
+
+    def deduct_inp_ranks_and_dtype(
+        self, out_abs_tensor: List[AbsTensor]
+    ) -> List[Tuple[int, DType]]:
+        return [(out_abs_tensor[0].ndims, out_abs_tensor[0].dtype)]
diff --git a/nnsmith/materialize/tensorflow/forward.py b/nnsmith/materialize/tensorflow/forward.py
new file mode 100644
index 00000000..ee1b63e0
--- /dev/null
+++ b/nnsmith/materialize/tensorflow/forward.py
@@ -0,0 +1,50 @@
+from functools import partial
+from typing import List, Type
+
+import tensorflow as tf  # type: ignore
+from keras import layers
+from multipledispatch import dispatch
+from tensorflow import keras
+
+from nnsmith.abstract.op import *
+from nnsmith.materialize import framework_operator_impl
+from nnsmith.materialize.tensorflow.dialect import Dense
+
+# core dialect + some future PyTorch-only Operators.
+# TF_REALIZABLE_OPS = FULL_OPERATOR_SETS["core"] + FULL_OPERATOR_SETS["tensorflow"]
+TF_REALIZABLE_OPS = [Add, Dense]
+ALL_TF_OPS: List[Type[AbsOpBase]] = []
+
+operator_impl = partial(framework_operator_impl, TF_REALIZABLE_OPS, ALL_TF_OPS)
+
+
+class StopFoldConst(tf.Module):
+    def __init__(self, data: tf.Tensor):
+        super().__init__()
+        self.data = tf.Variable(data, trainable=False)  # default variable is trainable
+
+    def __call__(self, training=None):
+        return self.data
+
+
+"""Implement TensorFlow forward Callables for operator classes"""
+
+
+@operator_impl(Constant)
+def forward_fn(op: Constant):
+    data = tf.cast(
+        tf.random.normal(op.abs_tensor.shape), op.abs_tensor.dtype.tensorflow()
+    )
+    return StopFoldConst(data)
+
+
+@operator_impl(Add)
+def forward_fn(op: Add):
+    return tf.add
+
+
+@operator_impl(Dense)
+def forward_fn(op: Dense):
+    return layers.Dense(
+        units=op.ofeat, dtype=op.input_like[0].dtype.tensorflow(), autocast=False
+    )
diff --git a/nnsmith/materialize/tensorflow/tfnet.py b/nnsmith/materialize/tensorflow/tfnet.py
new file mode 100644
index 00000000..5b95cc62
--- /dev/null
+++ b/nnsmith/materialize/tensorflow/tfnet.py
@@ -0,0 +1,90 @@
+from dataclasses import dataclass
+from typing import Callable, Dict, List, cast
+
+import tensorflow as tf
+
+from nnsmith.abstract.op import AbsOpBase, Input
+from nnsmith.error import SanityCheck
+from nnsmith.materialize import Schedule
+from nnsmith.materialize.tensorflow.forward import forward_fn
+
+
+@dataclass
+class Instr:
+    fwd_fn: Callable
+    inp_keys: List[int]
+    out_keys: List[int]
+
+
+TFNetOutDict = Dict[str, tf.Tensor]
+
+
+class TFNet(tf.Module):
+    """
+    Concrete TensorFlow Network
+    It only has minimal methods to be a TF network.
+    It only has minimal information "schedule" to do computation.
+    """
+
+    def __init__(
+        self,
+        schedule: Schedule,
+        verbose: bool = False,
+    ) -> None:
+        """Build a TensorFlow model from schedule
+
+        Args:
+            schedule (Schedule): minimal information for constructing a concrete graph.
+        """
+        super().__init__()
+        self.verbose = verbose
+        self.schedule: Schedule = schedule
+        self.mlist: List[Callable] = []
+        self.instructions: List[Instr] = []
+
+        for op, inp_keys, out_keys in self.schedule.instructions:
+            if not isinstance(op, Input):
+                op = cast(AbsOpBase, op)
+                fwd_fn = forward_fn(op)
+                SanityCheck.true(fwd_fn is not None, f"Bad implementation for {op}")
+                if not isinstance(op, tf.Module):
+                    self.mlist.append(fwd_fn)  # Add tf.Module to track its parameters
+                self.instructions.append(Instr(fwd_fn, inp_keys, out_keys))
+        # end for
+
+    @tf.function
+    def __call__(self, *args, **kwargs) -> TFNetOutDict:
+        if self.verbose:
+            mode = "Running Eagerly" if tf.executing_eagerly() else "Tracing"
+            print(f"{mode} with JIT config: {tf.config.optimizer.get_jit()}")
+
+        key2tensor: Dict[int, tf.Tensor] = {}
+        if len(args) == len(self.schedule.input_keys):
+            for i, key in enumerate(self.schedule.input_keys):
+                key2tensor[key] = args[i]
+        elif len(kwargs) == len(self.schedule.input_keys):
+            for i, key in enumerate(self.schedule.input_keys):
+                key2tensor[key] = kwargs[f"i{i}"]
+        else:
+            raise ValueError("Either user args only or kwargs only")
+
+        for i_instr, instr in enumerate(self.instructions):
+            # get inputs
+            inp_tensors = [key2tensor[key] for key in instr.inp_keys]
+
+            # forward
+            out_tensors = instr.fwd_fn(
+                *inp_tensors
+            )  # TODO Colin when it can return a list?
+            if not isinstance(out_tensors, list):
+                out_tensors = [out_tensors]
+
+            # store outputs
+            for i_out, out_key in enumerate(instr.out_keys):
+                key2tensor[out_key] = out_tensors[i_out]
+
+        # end for instructions
+        out_dict = {
+            f"o{i}": key2tensor[key] for i, key in enumerate(self.schedule.leaf_keys)
+        }
+        return out_dict
diff --git a/nnsmith/materialize/torch/__init__.py b/nnsmith/materialize/torch/__init__.py
index 341f56bd..7e9e78cc 100644
--- a/nnsmith/materialize/torch/__init__.py
+++ b/nnsmith/materialize/torch/__init__.py
@@ -1,11 +1,11 @@
-from typing import Dict, List, Type
 import pickle
+from os import PathLike
+from typing import Dict, List, Type
 
 import torch
 
-from nnsmith.graph_gen import Schedule
-from nnsmith.abstract.op import AbsTensor, AbsOpBase
-from nnsmith.materialize import Model, Oracle
+from nnsmith.abstract.op import AbsOpBase, AbsTensor
+from nnsmith.materialize import Model, Oracle, Schedule
 from nnsmith.materialize.torch.forward import ALL_TORCH_OPS
 from nnsmith.materialize.torch.input_gen import PracticalHybridSearch
 from nnsmith.materialize.torch.symbolnet import SymbolNet
@@ -58,7 +58,7 @@ def make_oracle(self) -> Oracle:
 
         return Oracle(input_dict, output_dict)
 
-    def dump(self, path: str):
+    def dump(self, path: PathLike):
         torch.save(self.torch_model.state_dict(), path)
         schedule_path = path.replace(
             TorchModel.name_prefix() + TorchModel.name_suffix(),
@@ -68,7 +68,7 @@ def dump(self, path: str):
             pickle.dump(self.torch_model.schedule, f)
 
     @staticmethod
-    def load(path: str) -> "TorchModel":
+    def load(path: PathLike) -> "TorchModel":
         ret = TorchModel()
         schedule_path = path.replace(
             TorchModel.name_prefix() + TorchModel.name_suffix(),
diff --git a/nnsmith/materialize/torch/dialect.py b/nnsmith/materialize/torch/dialect.py
index ea1d2eae..45b9aca9 100644
--- a/nnsmith/materialize/torch/dialect.py
+++ b/nnsmith/materialize/torch/dialect.py
@@ -1,8 +1,8 @@
 from typing import List, Tuple, Union
 
 from nnsmith.abstract.arith import *
-from nnsmith.abstract.op import int_from, UnaryOpBase, mark_materialize
 from nnsmith.abstract.dtype import DType
+from nnsmith.abstract.op import UnaryOpBase, int_from, mark_materialize
 from nnsmith.abstract.tensor import AbsTensor
 from nnsmith.error import ConstraintCheck
 
diff --git a/nnsmith/materialize/torch/forward.py b/nnsmith/materialize/torch/forward.py
index 02b9b7a0..d341b4c2 100644
--- a/nnsmith/materialize/torch/forward.py
+++ b/nnsmith/materialize/torch/forward.py
@@ -1,10 +1,13 @@
+from ast import operator
+from functools import partial
 from typing import Set, Type
-from multipledispatch import dispatch
 
 import torch
+from multipledispatch import dispatch
 
-from nnsmith.error import SanityCheck
 from nnsmith.abstract.op import *
+from nnsmith.error import SanityCheck
+from nnsmith.materialize import framework_operator_impl
 from nnsmith.materialize.torch.dialect import Linear
 
 # Implementation of operators.
@@ -13,25 +16,7 @@
 TORCH_REALIZABLE_OPS = FULL_OPERATOR_SETS["core"] + FULL_OPERATOR_SETS["torch"]
 ALL_TORCH_OPS: List[Type[AbsOpBase]] = []
 
-
-def operator_impl(op_type, *args, **kwargs):
-    SanityCheck.true(
-        issubclass(op_type, AbsOpBase),
-        f"Decorator operator_impl takes AbsOpBase subclass, but got {op_type}",
-    )
-    if op_type is not Constant:  # Constant comes from placeholder.
-        dispatchables = [
-            rtype for rtype in TORCH_REALIZABLE_OPS if issubclass(rtype, op_type)
-        ]
-        for rtype in dispatchables:
-            ALL_TORCH_OPS.append(rtype)
-
-        SanityCheck.true(
-            len(dispatchables) != 0,
-            f"Decorator operator_impl only take types decorated by `mark_realize`, but got {op_type}",
-        )
-    return dispatch(op_type, *args, **kwargs)
-
+operator_impl = partial(framework_operator_impl, TORCH_REALIZABLE_OPS, ALL_TORCH_OPS)
 
 # forward_fn:  forward
 
@@ -41,7 +26,7 @@ def __init__(self, data: torch.Tensor):
         super().__init__()
         self.dtype = data.dtype
         self.param = torch.nn.parameter.Parameter(
-            data, requires_grad=data.is_floating_point()
+            data, requires_grad=data.is_floating_point()  # TODO Colin trainable const?
         )
 
     @torch.no_grad()
diff --git a/nnsmith/materialize/torch/input_gen.py b/nnsmith/materialize/torch/input_gen.py
index e26d7930..fae104e0 100644
--- a/nnsmith/materialize/torch/input_gen.py
+++ b/nnsmith/materialize/torch/input_gen.py
@@ -1,12 +1,12 @@
 import time
-from typing import List, Tuple, Dict
 from abc import ABC, abstractmethod
+from typing import Dict, List, Tuple
 
-import torch
 import numpy as np
+import torch
 
-from nnsmith.materialize.torch.symbolnet import SymbolNet, random_tensor
 from nnsmith.abstract.op import DType
+from nnsmith.materialize.torch.symbolnet import SymbolNet, random_tensor
 
 
 class InputSearchBase(ABC):
diff --git a/nnsmith/materialize/torch/numeric.py b/nnsmith/materialize/torch/numeric.py
index e02dfa12..cf783075 100644
--- a/nnsmith/materialize/torch/numeric.py
+++ b/nnsmith/materialize/torch/numeric.py
@@ -1,5 +1,5 @@
-from multipledispatch import dispatch
 import torch
+from multipledispatch import dispatch
 
 from nnsmith.abstract.op import *
 
diff --git a/nnsmith/materialize/torch/proxy_grad.py b/nnsmith/materialize/torch/proxy_grad.py
index 86911a2f..ec4302f6 100644
--- a/nnsmith/materialize/torch/proxy_grad.py
+++ b/nnsmith/materialize/torch/proxy_grad.py
@@ -1,5 +1,5 @@
-from multipledispatch import dispatch
 import torch
+from multipledispatch import dispatch
 
 from nnsmith.abstract.op import *
 
diff --git a/nnsmith/materialize/torch/symbolnet.py b/nnsmith/materialize/torch/symbolnet.py
index c610b401..0ba61e28 100644
--- a/nnsmith/materialize/torch/symbolnet.py
+++ b/nnsmith/materialize/torch/symbolnet.py
@@ -1,17 +1,17 @@
 import os
 import time
 import warnings
-from typing import Optional, Dict
+from typing import Dict, Optional
 
 import torch
 from torch import nn
 
+from nnsmith.abstract.op import Input
+from nnsmith.error import ConstraintCheck, ConstraintError, SanityCheck
+from nnsmith.materialize import Schedule
+from nnsmith.materialize.torch.forward import forward_fn
 from nnsmith.materialize.torch.numeric import loss_fn, numeric_valid
 from nnsmith.materialize.torch.proxy_grad import proxy_fn
-from nnsmith.materialize.torch.forward import forward_fn
-from nnsmith.graph_gen import Schedule
-from nnsmith.error import ConstraintCheck, ConstraintError, SanityCheck
-from nnsmith.abstract.op import Input
 
 __MB_LIM__ = 6 * 1024
 __INPUT_FOUND_NAN_MSG__ = "[NaN] in model inputs!"
@@ -75,7 +75,9 @@ def __init__(
                 if loss_fn.dispatch(type(op)):
                     self.n_vulnerable_op += 1
 
-                self.instructions.append((torch_fn, inputs, outputs, op))
+                self.instructions.append(
+                    (torch_fn, inputs, outputs, op)
+                )  # TODO Colin op is redundant?
 
         # the order follows `input_keys`
         self.input_map = {
diff --git a/nnsmith/narrow_spec.py b/nnsmith/narrow_spec.py
index 19b8bad1..0c43b960 100644
--- a/nnsmith/narrow_spec.py
+++ b/nnsmith/narrow_spec.py
@@ -18,21 +18,20 @@
 
 
 import os
-from os import PathLike
-from dataclasses import dataclass
-from typing import List, Dict, Type
 from copy import deepcopy
+from dataclasses import dataclass
+from os import PathLike
+from typing import Dict, List, Type
 
-from omegaconf import OmegaConf
-from appdirs import user_cache_dir
 import z3
+from appdirs import user_cache_dir
+from omegaconf import OmegaConf
 
-from nnsmith.abstract.op import Input, Constant, AbsOpBase, AbsTensor, concretize_op
 from nnsmith.abstract.dtype import DType
+from nnsmith.abstract.op import AbsOpBase, AbsTensor, Constant, Input, concretize_op
 from nnsmith.backends import BackendFactory
-from nnsmith.graph_gen import Instruction, Schedule
-from nnsmith.materialize import Model, TestCase
-from nnsmith.cli.util import succ_print, fail_print, note_print
+from nnsmith.cli.util import fail_print, note_print, succ_print
+from nnsmith.materialize import Instruction, Model, Schedule, TestCase
 
 NNSMITH_CACHE_DIR = user_cache_dir("nnsmith")
 
diff --git a/nnsmith/util.py b/nnsmith/util.py
index 411bc817..729b2c9f 100644
--- a/nnsmith/util.py
+++ b/nnsmith/util.py
@@ -1,11 +1,10 @@
 import os
-import shutil
 import random
+import shutil
 from typing import Callable, Dict, List
 
 import numpy as np
 
-
 SEED_SETTERS = {
     "random": random.seed,
     "numpy": np.random.seed,
diff --git a/requirements/core.txt b/requirements/core.txt
index 4f5350e2..58130337 100644
--- a/requirements/core.txt
+++ b/requirements/core.txt
@@ -6,3 +6,4 @@ termcolor
 multipledispatch
 gitpython
 hydra-core
+appdirs
diff --git a/requirements/install_all.sh b/requirements/install_all.sh
new file mode 100644
index 00000000..3ed9e585
--- /dev/null
+++ b/requirements/install_all.sh
@@ -0,0 +1,30 @@
+#!/bin/bash
+
+set -ex
+
+cd ./requirements
+
+# detect cuda
+# check if nvcc can work and return true or false
+check_nvcc() {
+    nvcc -V > /dev/null 2>&1
+    return $has_cuda
+}
+
+for dep_file in $(find . -name '*.txt')
+do
+  if [ -f "$dep_file" ]; then
+      # skip tensorrt if cuda is not available
+      if [[ "$dep_file" == *"tensorrt"* ]] && ! check_cuda; then
+          echo "Skipping tensorrt"
+          continue
+      fi
+
+      UPGRADE_PACKAGES=""
+      if [[ "$dep_file" == *"sys"* ]]; then
+          # Files under sys should be nightly releases
+          UPGRADE_PACKAGES="--upgrade"
+      fi
+      pip install -r "$dep_file" $UPGRADE_PACKAGES
+  fi
+done
diff --git a/requirements/sys/tensorflow.txt b/requirements/sys/tensorflow.txt
new file mode 100644
index 00000000..dd3d4130
--- /dev/null
+++ b/requirements/sys/tensorflow.txt
@@ -0,0 +1,3 @@
+tf-nightly
+# NOTE: https://www.tensorflow.org/install
+# https://www.tensorflow.org/install/pip


Test patch:
diff --git a/nnsmith/cli/dtype_test.py b/nnsmith/cli/dtype_test.py
index 68e32cbd..2a8d502a 100644
--- a/nnsmith/cli/dtype_test.py
+++ b/nnsmith/cli/dtype_test.py
@@ -1,8 +1,8 @@
 import hydra
 from omegaconf import DictConfig
 
-from nnsmith.materialize import Model
 from nnsmith.backends import BackendFactory
+from nnsmith.materialize import Model
 from nnsmith.narrow_spec import load_topset_from_auto_cache
 
 
diff --git a/nnsmith/difftest.py b/nnsmith/difftest.py
index 64e336e2..4016d6cc 100644
--- a/nnsmith/difftest.py
+++ b/nnsmith/difftest.py
@@ -1,4 +1,5 @@
 from typing import Dict
+
 import numpy as np
 from numpy import testing
 
diff --git a/tests/test_native_pickable.py b/tests/core/test_native_pickable.py
similarity index 100%
rename from tests/test_native_pickable.py
rename to tests/core/test_native_pickable.py
index f3b126c7..20099101 100644
--- a/tests/test_native_pickable.py
+++ b/tests/core/test_native_pickable.py
@@ -1,10 +1,10 @@
-import pytest
-
 import pickle
 
+import pytest
+
 from nnsmith.abstract.dtype import DType
-from nnsmith.abstract.tensor import AbsTensor
 from nnsmith.abstract.op import MaxPool2d
+from nnsmith.abstract.tensor import AbsTensor
 
 
 def test_dtype_picklable():
diff --git a/tests/tensorflow/test_tflite_backend.py b/tests/tensorflow/test_tflite_backend.py
new file mode 100644
index 00000000..05089a45
--- /dev/null
+++ b/tests/tensorflow/test_tflite_backend.py
@@ -0,0 +1,43 @@
+import pytest
+import tensorflow as tf
+
+from nnsmith.backends.tflite import TFLiteFactory
+from nnsmith.graph_gen import concretize_graph, random_model_gen
+from nnsmith.materialize import Schedule, TestCase
+from nnsmith.materialize.tensorflow import TFModel
+
+TestCase.__test__ = False  # supress PyTest warning
+
+
+def test_synthesized_tflite_model(tmp_path):
+    d = tmp_path / "test_tflite"
+    d.mkdir()
+
+    # TODO(@ganler): do dtype first.
+    gen = random_model_gen(
+        opset=TFModel.operators(),
+        init_rank=4,
+        seed=23132,
+        max_nodes=4,
+    )  # One op should not be easily wrong... I guess.
+
+    fixed_graph, concrete_abstensors = concretize_graph(
+        gen.abstract_graph, gen.tensor_dataflow, gen.get_solutions()
+    )
+
+    schedule = Schedule.init(fixed_graph, concrete_abstensors)
+
+    model = TFModel.from_schedule(schedule)
+
+    # model.refine_weights()  # either random generated or gradient-based.
+    oracle = model.make_oracle()
+
+    testcase = TestCase(model, oracle)
+    testcase.dump(root_folder=d)
+
+    assert (
+        TFLiteFactory(
+            device="cpu", optmax=False, catch_process_crash=False
+        ).verify_testcase(testcase)
+        is None
+    )
diff --git a/tests/test_dump_load.py b/tests/torch/test_dump_load.py
similarity index 92%
rename from tests/test_dump_load.py
rename to tests/torch/test_dump_load.py
index c7f1cb81..b089162a 100644
--- a/tests/test_dump_load.py
+++ b/tests/torch/test_dump_load.py
@@ -1,11 +1,10 @@
-import pytest
-
 import numpy as np
+import pytest
 import torch
 
-from nnsmith.materialize import TestCase
+from nnsmith.graph_gen import concretize_graph, random_model_gen
+from nnsmith.materialize import Schedule, TestCase
 from nnsmith.materialize.onnx import ONNXModel
-from nnsmith.graph_gen import random_model_gen, concretize_graph, make_schedule
 
 TestCase.__test__ = False  # supress PyTest warning
 
@@ -25,7 +24,7 @@ def test_onnx_load_dump(tmp_path):
         gen.abstract_graph, gen.tensor_dataflow, gen.get_solutions()
     )
 
-    schedule = make_schedule(fixed_graph, concrete_abstensors)
+    schedule = Schedule.init(fixed_graph, concrete_abstensors)
 
     model = ONNXModel.from_schedule(schedule)
 
@@ -71,7 +70,7 @@ def test_bug_report_load_dump(tmp_path):
         gen.abstract_graph, gen.tensor_dataflow, gen.get_solutions()
     )
 
-    schedule = make_schedule(fixed_graph, concrete_abstensors)
+    schedule = Schedule.init(fixed_graph, concrete_abstensors)
 
     model = ONNXModel.from_schedule(schedule)
 
diff --git a/tests/test_ort_backend.py b/tests/torch/test_ort_backend.py
similarity index 90%
rename from tests/test_ort_backend.py
rename to tests/torch/test_ort_backend.py
index 534550ce..c6d492b3 100644
--- a/tests/test_ort_backend.py
+++ b/tests/torch/test_ort_backend.py
@@ -1,9 +1,9 @@
 import pytest
 
 from nnsmith.abstract.dtype import DType
-from nnsmith.materialize import TestCase, Model
-from nnsmith.graph_gen import random_model_gen, concretize_graph, make_schedule
 from nnsmith.backends import BackendFactory
+from nnsmith.graph_gen import concretize_graph, random_model_gen
+from nnsmith.materialize import Model, Schedule, TestCase
 from nnsmith.narrow_spec import load_topset_from_auto_cache
 
 TestCase.__test__ = False  # supress PyTest warning
@@ -26,7 +26,7 @@ def test_synthesized_onnx_model(tmp_path):
         gen.abstract_graph, gen.tensor_dataflow, gen.get_solutions()
     )
 
-    schedule = make_schedule(fixed_graph, concrete_abstensors)
+    schedule = Schedule.init(fixed_graph, concrete_abstensors)
 
     model = ONNXModel.from_schedule(schedule)
 
diff --git a/tests/test_trt_backend.py b/tests/torch/test_trt_backend.py
similarity index 90%
rename from tests/test_trt_backend.py
rename to tests/torch/test_trt_backend.py
index 34f7d712..947e5df3 100644
--- a/tests/test_trt_backend.py
+++ b/tests/torch/test_trt_backend.py
@@ -1,6 +1,5 @@
-import pytest
-
 import GPUtil
+import pytest
 
 if not GPUtil.getAvailable():
     pytest.skip(
@@ -8,10 +7,9 @@
     )
 
 from nnsmith.abstract.dtype import DType
-from nnsmith.materialize import TestCase
-from nnsmith.materialize import Model
-from nnsmith.graph_gen import random_model_gen, concretize_graph, make_schedule
 from nnsmith.backends import BackendFactory
+from nnsmith.graph_gen import concretize_graph, random_model_gen
+from nnsmith.materialize import Model, Schedule, TestCase
 from nnsmith.narrow_spec import load_topset_from_auto_cache
 
 TestCase.__test__ = False  # supress PyTest warning
@@ -35,7 +33,7 @@ def test_synthesized_onnx_model(tmp_path):
         gen.abstract_graph, gen.tensor_dataflow, gen.get_solutions()
     )
 
-    schedule = make_schedule(fixed_graph, concrete_abstensors)
+    schedule = Schedule.init(fixed_graph, concrete_abstensors)
 
     model = ONNXModel.from_schedule(schedule)
 
diff --git a/tests/test_tvm_backend.py b/tests/torch/test_tvm_backend.py
similarity index 91%
rename from tests/test_tvm_backend.py
rename to tests/torch/test_tvm_backend.py
index 3894267f..1bd41f81 100644
--- a/tests/test_tvm_backend.py
+++ b/tests/torch/test_tvm_backend.py
@@ -1,11 +1,10 @@
 import pytest
-
 import tvm
 
 from nnsmith.abstract.dtype import DType
-from nnsmith.materialize import TestCase, Model
-from nnsmith.graph_gen import random_model_gen, concretize_graph, make_schedule
 from nnsmith.backends import BackendFactory
+from nnsmith.graph_gen import concretize_graph, random_model_gen
+from nnsmith.materialize import Model, Schedule, TestCase
 from nnsmith.narrow_spec import load_topset_from_auto_cache
 
 TestCase.__test__ = False  # supress PyTest warning
@@ -28,7 +27,7 @@ def test_synthesized_onnx_model(tmp_path):
         gen.abstract_graph, gen.tensor_dataflow, gen.get_solutions()
     )
 
-    schedule = make_schedule(fixed_graph, concrete_abstensors)
+    schedule = Schedule.init(fixed_graph, concrete_abstensors)
 
     model = ONNXModel.from_schedule(schedule)