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

Issue:
Does NNSmith support fp16?
Hi. Thanks for developing this amazing project!

Just curious if nnsmith supports generating tests with dtype float16? Directly adding float16 to `abstract/dtype.py` seems not working.


Discussion / hints:
We had that in an experimental (private) repo. I am happy to upstream the feature you mentioned later today.
Thanks! Looking forward to it.

Fix patch:
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
index 4b36860..504fe4c 100644
--- a/.github/workflows/ci.yaml
+++ b/.github/workflows/ci.yaml
@@ -38,7 +38,6 @@ jobs:
       - name: Test TensorFlow
         run: |
           pip install -r requirements/sys/tensorflow.txt --pre --upgrade
-          pip install -r requirements/sys/iree.txt --pre --upgrade
           pytest tests/tensorflow --log-cli-level=DEBUG
           yes | python nnsmith/cli/model_gen.py model.type=tensorflow mgen.method=symbolic
           python nnsmith/cli/model_exec.py model.type=tensorflow backend.type=xla model.path=nnsmith_output/model/
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 8bebfcc..383ac67 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -1,6 +1,6 @@
 repos:
 -   repo: https://github.com/pycqa/isort
-    rev: 5.10.1
+    rev: 5.12.0
     hooks:
     -   id: isort
         name: isort (python)
diff --git a/README.md b/README.md
index 1e99d1b..46ccbfb 100644
--- a/README.md
+++ b/README.md
@@ -15,7 +15,7 @@
 | Model\Engine | [TVM](https://github.com/apache/tvm) | [ORT](https://github.com/microsoft/onnxruntime) | [TensorRT](https://github.com/NVIDIA/TensorRT) | [TFLite](https://www.tensorflow.org/lite) | [XLA](https://www.tensorflow.org/xla) | [Torch-JIT](https://pytorch.org/docs/stable/jit.html) |
 | ------------ | ------------------------------------ | ----------------------------------------------- | ---------------------------------------------- | ----------------------------------------- | ------------------------------------- | ----------------------------------------------------- |
 | ONNX         | ✅                                    | ✅                                               | ✅                                              |                                           |                                       |                                                       |
-| TensorFlow   | 🔨                                    |                                                 |                                                | ⚠️                                         | ⚠️                                     |                                                       |
+| TensorFlow   | 🔨                                    |                                                 |                                                | ✅                                         | ✅                                     |                                                       |
 | PyTorch      | 🔨                                    | 🔨                                               |                                                |                                           |                                       | 🔨                                                     |
 
 
@@ -27,24 +27,24 @@
 
 ## Setup
 
-**Install latest stable release:**
+**Install latest code (GitHub HEAD):**
 
 ```shell
-pip install "nnsmith[torch,onnx]" --upgrade
+pip install "git+https://github.com/ise-uiuc/nnsmith@main#egg=nnsmith[torch,onnx]" --upgrade
+# [optional] add more front- and back-ends such as [tf] and [tvm,ort,xla,...] in "[...]"
 ```
 
-<details><summary><b>Install GitHub HEAD: </b> <i>[click to expand]</i></summary>
+<details><summary><b>Install latest stable release: </b> <i>[expand]</i></summary>
 <div>
 
 ```shell
-pip install "git+https://github.com/ise-uiuc/nnsmith@main#egg=nnsmith[torch,onnx]" --upgrade
-# or pip install "git+ssh://git@github.com/ise-uiuc/nnsmith@main#egg=nnsmith[torch,onnx]" --upgrade
+pip install "nnsmith[torch,onnx]" --upgrade
 ```
 
 </div>
 </details>
 
-<details><summary><b>Install latest pre-release: </b> <i>[click to expand]</i></summary>
+<details><summary><b>Install latest pre-release: </b> <i>[expand]</i></summary>
 <div>
 
 ```shell
@@ -60,7 +60,7 @@ pip install "nnsmith[torch,onnx]"                     \
 
 ## Quick Start
 
-<details><summary><b>Setting up graphviz for debugging</b> <i>[click to expand]</i></summary>
+<details><summary><b>Setting up graphviz for debugging</b> <i>[expand]</i></summary>
 <div>
 
 Graphviz provides `dot` for visualizing graphs in nice pictures. But it needs to be installed via the following methods:
@@ -92,7 +92,7 @@ See other commands under [`doc/cli`](doc/cli.md). We use [hydra](https://hydra.c
 - `pip install --upgrade --pre -r requirements/sys/[system].txt` to allow generating and running specific frameworks;
   -  **Why "--upgrade --pre"?** In fact, all the sources under `requirements/sys/` are nightly release (except tvm) as we want to "save the world" by catching new bugs;
 
-<details><summary><b>Pre-commits</b> <i>[click to expand]</i></summary>
+<details><summary><b>Pre-commits</b> <i>[expand]</i></summary>
 <div>
 
 You can use `pre-commit` to simpify development:
@@ -104,7 +104,7 @@ You can use `pre-commit` to simpify development:
 </div>
 </details>
 
-<details><summary><b>Local development</b> <i>[click to expand]</i></summary>
+<details><summary><b>Local development</b> <i>[expand]</i></summary>
 <div>
 
 - Develop locally by setting `export PYTHONPATH=$PYTHONPATH:$(pwd)` (`pwd` should be this git folder.)
@@ -113,7 +113,7 @@ You can use `pre-commit` to simpify development:
 </div>
 </details>
 
-<details><summary><b>Simplify the code</b> <i>[click to expand]</i></summary>
+<details><summary><b>Simplify the code</b> <i>[expand]</i></summary>
 <div>
 
 *Simplicity is prerequisite for reliability.* --Edsger W. Dijkstra
@@ -123,7 +123,7 @@ We want **code simplicity**: keeping minimal dependencies and focusing on a smal
 </div>
 </details>
 
-<details><summary><b>Test before commit</b> <i>[click to expand]</i></summary>
+<details><summary><b>Test before commit</b> <i>[expand]</i></summary>
 <div>
 
 ```shell
@@ -144,15 +144,16 @@ pytest tests/tensorflow -s
 
 ## Paper
 
-<details><summary><b>ASPLOS'23 | NNSmith: Generating Diverse and Valid Test Cases for Deep Learning Compilers.</b> <i>[click to expand citation]</i></summary>
+<details><summary><b>NNSmith: Generating Diverse and Valid Test Cases for Deep Learning Compilers.</b> <i>[expand citation]</i></summary>
 <div>
 
 ```bibtex
-@article{liu2022finding,
-  title={Finding Deep-Learning Compilation Bugs with NNSmith},
+@inproceedings{liu2023nnsmith,
+  title={Nnsmith: Generating diverse and valid test cases for deep learning compilers},
   author={Liu, Jiawei and Lin, Jinkun and Ruffy, Fabian and Tan, Cheng and Li, Jinyang and Panda, Aurojit and Zhang, Lingming},
-  journal={arXiv preprint arXiv:2207.13066},
-  year={2022}
+  booktitle={Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2},
+  pages={530--543},
+  year={2023}
 }
 ```
 
@@ -160,6 +161,7 @@ pytest tests/tensorflow -s
 </details>
 
 <p align="center">
+    <a href="https://dl.acm.org/doi/10.1145/3575693.3575707"><img src="https://img.shields.io/badge/Paper-ASPLOS'23-a55fed.svg"></a>
     <a href="https://arxiv.org/abs/2207.13066"><img src="https://img.shields.io/badge/arXiv-2207.13066-b31b1b.svg"></a>
     <a href="http://nnsmith-asplos.rtfd.io/"><img src="https://img.shields.io/badge/artifact-doc-black.svg"></a>
     <a href="https://github.com/ganler/nnsmith-asplos-artifact"><img src="https://img.shields.io/badge/artifact-git-black.svg"></a>
diff --git a/doc/cli.md b/doc/cli.md
index 5f3dbd8..bbddba6 100644
--- a/doc/cli.md
+++ b/doc/cli.md
@@ -19,7 +19,8 @@ and backends:
 - `tvm`
 - `ort`: ONNXRuntime
 - `trt`: TensorRT
-- `iree`
+- `xla`: XLA
+- `tflite`: TFLite
 
 Meanwhile, the backend of `xla` and `tflite` is installed as part of TensorFlow.
 
diff --git a/nnsmith/abstract/dtype.py b/nnsmith/abstract/dtype.py
index 5236bb6..99a0caf 100644
--- a/nnsmith/abstract/dtype.py
+++ b/nnsmith/abstract/dtype.py
@@ -9,6 +9,9 @@ class DType(Enum):
     float32 = "float32"
     float64 = "float64"
     uint8 = "uint8"  # Support quantized models.
+    uint16 = "uint16"
+    uint32 = "uint32"
+    uint64 = "uint64"
     int8 = "int8"
     int16 = "int16"
     int32 = "int32"
@@ -31,6 +34,9 @@ def short(self) -> str:
             DType.float32: "f32",
             DType.float64: "f64",
             DType.uint8: "u8",
+            DType.uint16: "u16",
+            DType.uint32: "u32",
+            DType.uint64: "u64",
             DType.int8: "i8",
             DType.int16: "i16",
             DType.int32: "i32",
@@ -42,6 +48,7 @@ def short(self) -> str:
 
     @staticmethod
     def is_float(dtype):  # Don't use string. Make it well-formed.
+        assert isinstance(dtype, DType)
         return dtype in [DType.float32, DType.float64]
 
     @staticmethod
@@ -60,6 +67,9 @@ def from_str(s):
             "float32": DType.float32,
             "float64": DType.float64,
             "uint8": DType.uint8,
+            "uint16": DType.uint16,
+            "uint32": DType.uint32,
+            "uint64": DType.uint64,
             "int8": DType.int8,
             "int16": DType.int16,
             "int32": DType.int32,
@@ -75,6 +85,10 @@ def numpy(self):
             DType.float32: np.float32,
             DType.float64: np.float64,
             DType.uint8: np.uint8,
+            DType.uint8: np.uint8,
+            DType.uint16: np.uint16,
+            DType.uint32: np.uint32,
+            DType.uint64: np.uint64,
             DType.int8: np.int8,
             DType.int16: np.int16,
             DType.int32: np.int32,
@@ -84,7 +98,6 @@ def numpy(self):
             DType.bool: np.bool_,
         }[self]
 
-    # TODO(@ganler): put "torchization" in a separate file.
     def torch(self) -> "torch.dtype":
         import torch
 
@@ -93,6 +106,7 @@ def torch(self) -> "torch.dtype":
             DType.float32: torch.float32,
             DType.float64: torch.float64,
             DType.uint8: torch.uint8,
+            # PyTorch does not support other unsigned int types: https://github.com/pytorch/pytorch/issues/58734
             DType.int8: torch.int8,
             DType.int16: torch.int16,
             DType.int32: torch.int32,
@@ -128,6 +142,9 @@ def tensorflow(self) -> "tf.Dtype":
             DType.float32: tf.float32,
             DType.float64: tf.float64,
             DType.uint8: tf.uint8,
+            DType.uint16: tf.uint16,
+            DType.uint32: tf.uint32,
+            DType.uint64: tf.uint64,
             DType.int8: tf.int8,
             DType.int16: tf.int16,
             DType.int32: tf.int32,
@@ -146,6 +163,9 @@ def from_tensorflow(dtype) -> "DType":
             tf.float32: DType.float32,
             tf.float64: DType.float64,
             tf.uint8: DType.uint8,
+            tf.uint16: DType.uint16,
+            tf.uint32: DType.uint32,
+            tf.uint64: DType.uint64,
             tf.int8: DType.int8,
             tf.int16: DType.int16,
             tf.int32: DType.int32,
@@ -161,6 +181,9 @@ def sizeof(self) -> int:
             DType.float32: 4,
             DType.float64: 8,
             DType.uint8: 1,
+            DType.uint16: 2,
+            DType.uint32: 4,
+            DType.uint64: 8,
             DType.int8: 1,
             DType.int16: 2,
             DType.int32: 4,
@@ -175,14 +198,14 @@ def sizeof(self) -> int:
 # "DTYPE_GEN_ALL" is surely a subset of all types but it is
 # used to conservatively to avoid unsupported data types while
 # applying nnsmith to various frameworks.
-DTYPE_ALL = [dt for dt in DType]
-DTYPE_GEN_ALL = [
-    DType.float32,
-    DType.float64,
+DTYPE_GEN_FLOATS = [DType.float16, DType.float32, DType.float64]
+DTYPE_GEN_INTS = [
+    DType.int8,
+    DType.int16,
     DType.int32,
     DType.int64,
-    DType.bool,
+    DType.uint8,
 ]
+DTYPE_GEN_COMPLEX = [DType.complex64, DType.complex128]
+DTYPE_GEN_ALL = DTYPE_GEN_FLOATS + DTYPE_GEN_INTS + DTYPE_GEN_COMPLEX
 DTYPE_GEN_NON_BOOL = [dtype for dtype in DTYPE_GEN_ALL if dtype != DType.bool]
-DTYPE_GEN_FLOATS = [DType.float32, DType.float64]
-DTYPE_GEN_INTS = [DType.int32, DType.int64]
diff --git a/nnsmith/abstract/op.py b/nnsmith/abstract/op.py
index de6ce46..628aacc 100644
--- a/nnsmith/abstract/op.py
+++ b/nnsmith/abstract/op.py
@@ -11,6 +11,7 @@
 from nnsmith.abstract.arith import *
 from nnsmith.abstract.dtype import (
     DTYPE_GEN_ALL,
+    DTYPE_GEN_COMPLEX,
     DTYPE_GEN_FLOATS,
     DTYPE_GEN_INTS,
     DTYPE_GEN_NON_BOOL,
@@ -23,24 +24,7 @@
 # There are following types of constraints at this point:
 # 1. Shape variables must be greater than 0;
 # 2. Shape variables must avoid devision by 0;
-# 3. Intra-input shape constraints; e.g., add(x, y) where x.shape() must be equal to y.shape();
-# 4. Extra constraints introduced by individual operators;
-
-FLOPS_LIM = os.getenv("NNSMITH_FLOPS_LIM", "auto")
-if FLOPS_LIM == "auto":  # use predefined value
-    FLOPS_LIM = 2**30
-elif FLOPS_LIM == "off":
-    FLOPS_LIM = None
-else:
-    FLOPS_LIM = float(FLOPS_LIM)
-
-# control wheter to model FLOPS in z3 too. If not, we will check it after model is concretized.
-Z3_CONS_FLOPS = os.getenv("NNSMITH_Z3_CONS_FLOPS", "on")
-assert Z3_CONS_FLOPS in [
-    "on",
-    "off",
-], "NNSMITH_Z3_CONS_FLOPS must be either 'on' or 'off'"
-Z3_CONS_FLOPS = Z3_CONS_FLOPS == "on"
+# 3. Extra constraints introduced by individual operators;
 
 
 __MIN_RANK__ = 0
@@ -151,6 +135,10 @@ def z3_bcast(
     )
 
 
+def int_bcast(x: int, *args: int):
+    return x if (x != 1 or len(args) == 0) else int_bcast(*args)
+
+
 def broadcast_shapes(
     *shapes: List[Union[z3.ExprRef, int]]
 ) -> List[Union[z3.ExprRef, int]]:
@@ -159,15 +147,15 @@ def broadcast_shapes(
     if len(shapes) == 1:
         return shapes[0]
     max_dim = max(map(lambda x: len(x), shapes))
-    max_shape = [None] * (max_dim)
+    out_shape = [None] * max_dim
     for j in range(max_dim):
         i = -j - 1
         args_dim_sz = [_prepend_to(x, max_dim)[i] for x in shapes]
         if any(isinstance(s, z3.ExprRef) for s in args_dim_sz):
-            max_shape[i] = z3.simplify(z3_bcast(*args_dim_sz))
+            out_shape[i] = z3.simplify(z3_bcast(*args_dim_sz))
         else:
-            max_shape[i] = max(*args_dim_sz)
-    return max_shape
+            out_shape[i] = int_bcast(*args_dim_sz)
+    return out_shape
 
 
 def broadcast_cons(*shapes: List[Union[z3.ExprRef, int]]) -> List[z3.BoolRef]:
@@ -724,7 +712,6 @@ class LeakyReLU(ElementWiseUnaryOp):
     def __init__(self):
         """See https://pytorch.org/docs/stable/generated/torch.nn.LeakyReLU.html"""
         super().__init__()
-        self.negative_slope = 0.01
 
 
 @mark_materialize("core")
@@ -781,7 +768,19 @@ class Atan(TrigonometricOp):
 
 @mark_materialize("core")
 class Abs(ElementWiseUnaryOp):
-    in_dtypes = [(i,) for i in DTYPE_GEN_NON_BOOL]
+    in_dtypes = [(i,) for i in DTYPE_GEN_NON_BOOL if i != DType.uint8]
+    in_dtypes = [
+        (i,)
+        for i in DTYPE_GEN_NON_BOOL
+        if i not in {DType.complex64, DType.complex128, DType.uint8}
+    ]
+
+    def type_transfer(self, input_shapes: List[AbsTensor]) -> List[AbsTensor]:
+        if input_shapes[0].dtype == DType.complex64:
+            return [AbsTensor(input_shapes[0].shape, DType.float32)]
+        elif input_shapes[0].dtype == DType.complex128:
+            return [AbsTensor(input_shapes[0].shape, DType.float64)]
+        return super().type_transfer(input_shapes)
 
 
 @mark_materialize("core")
@@ -927,9 +926,6 @@ def requires(self, input_shapes):
             [nnsmith_gt(v, 0) for v in input_shapes[0].shape]
         )  # dim cannot be 0 for maxpool
 
-        # limit FLOPS
-        if Z3_CONS_FLOPS:
-            cons.append(nnsmith_le(self.flops(input_shapes), FLOPS_LIM))
         for c in cons:
             ret.append(c)
         return ret
@@ -1002,11 +998,6 @@ def _get_attrs(self, ndims):
             self.extra_attrs["region"] = random.choice(["left", "mid", "right"])
             if random.uniform(0, 1) < 0.1:
                 # torch exporter does not support start=INT_MIN
-                # if random.uniform(0, 1) < 0.5:
-                #     # because pytorch only supports forward slicing,
-                #     # start cannot be INT_MAX, otherwise it slices empty tensor
-                #     self.start = self.INT_MIN
-                # else:
                 self.end = self.INT_MAX
         return self.extra_attrs["axis"]
 
@@ -1141,6 +1132,8 @@ def __init__(self, *padding_list):
 @mark_materialize("core")
 class ReflectPad(Pad):
     num_var_param = _pad_num_var_param(2, max=6)
+    in_dtypes = [(i,) for i in DTYPE_GEN_FLOATS if i != DType.float16]
+    out_dtypes = [(i,) for i in DTYPE_GEN_FLOATS if i != DType.float16]
 
     def __init__(self, *padding_list):
         super().__init__(padding_list, "reflect")
@@ -1441,9 +1434,6 @@ def requires(self, input_shapes):
         )
         # not too extream to avoid torch exporter issue
         cons.append(nnsmith_le(self.padding, 255))
-        # limit FLOPS
-        if Z3_CONS_FLOPS:
-            cons.append(nnsmith_le(self.flops(input_shapes), FLOPS_LIM))
         return cons
 
     def flops(self, input_shapes):
@@ -1550,7 +1540,6 @@ def type_transfer(self, input_shapes: List[AbsTensor]) -> List[AbsTensor]:
         auto_dim = -1
         accum = 1
         for i, v in enumerate(self.target_shape):
-            # TODO: What to do about bitvectors here?
             if v == -1:
                 SanityCheck.eq(auto_dim, -1)
                 auto_dim = i
@@ -1629,7 +1618,6 @@ class Transpose(UnaryOpBase):
     in_dtypes = [(i,) for i in DTYPE_GEN_ALL]
 
     def __init__(self):
-        """See https://pytorch.org/docs/stable/generated/torch.transpose.html"""
         super().__init__()
         self.inp_ranks = [rank_from(2)]
         self.out_ranks = [rank_from(2)]
@@ -1666,9 +1654,6 @@ def deduct_inp_ranks_and_dtype(
         return [(out_abs_tensor[0].ndims, out_abs_tensor[0].dtype)]
 
 
-# Sum, Min, Max, Mean, ArgMin, ArgMax, Squeeze, Size
-
-
 class InterpBase(UnaryOpBase):
     num_var_param = rank_range(1, 3)
 
@@ -1817,7 +1802,7 @@ class ReduceProd(ReduceBase):
 
 @mark_materialize("core")
 class ArgMin(ReduceBase):
-    in_dtypes = [(i,) for i in DTYPE_GEN_NON_BOOL]
+    in_dtypes = [(i,) for i in DTYPE_GEN_NON_BOOL if i not in DTYPE_GEN_COMPLEX]
     out_dtypes = [(DType.int64,)]
     _reduce_out_dtype = DType.int64
 
@@ -1831,7 +1816,7 @@ def deduct_inp_ranks_and_dtype(
 
 @mark_materialize("core")
 class ArgMax(ReduceBase):
-    in_dtypes = [(i,) for i in DTYPE_GEN_NON_BOOL]
+    in_dtypes = [(i,) for i in DTYPE_GEN_NON_BOOL if i not in DTYPE_GEN_COMPLEX]
     out_dtypes = [(DType.int64,)]
     _reduce_out_dtype = DType.int64
 
@@ -2041,8 +2026,8 @@ def __init__(self):
 
 @mark_materialize("core")
 class MatMul(BinaryOpBase):
-    in_dtypes = [(i, i) for i in DTYPE_GEN_NON_BOOL]
-    out_dtypes = [(i,) for i in DTYPE_GEN_NON_BOOL]
+    in_dtypes = [(i, i) for i in DTYPE_GEN_NON_BOOL if i not in DTYPE_GEN_COMPLEX]
+    out_dtypes = [(i,) for i in DTYPE_GEN_NON_BOOL if i not in DTYPE_GEN_COMPLEX]
 
     def __init__(self):
         super().__init__()
diff --git a/nnsmith/backends/factory.py b/nnsmith/backends/factory.py
index 4b36ddf..ae7121b 100644
--- a/nnsmith/backends/factory.py
+++ b/nnsmith/backends/factory.py
@@ -310,7 +310,7 @@ def make_testcase(
     def init(name, target="cpu", optmax=True, **kwargs):
         if name is None:
             raise ValueError(
-                "Backend type cannot be None. Specify via `backend.type=[onnxruntime | tvm | tensorrt | tflite | xla | iree]`"
+                "Backend type cannot be None. Specify via `backend.type=[onnxruntime | tvm | tensorrt | tflite | xla ]`"
             )
 
         if target == "gpu":
@@ -358,13 +358,5 @@ def init(name, target="cpu", optmax=True, **kwargs):
                 optmax=optmax,
                 **kwargs,
             )
-        elif name == "iree":
-            from nnsmith.backends.iree import IREEFactory
-
-            return IREEFactory(
-                target=target,
-                optmax=optmax,
-                **kwargs,
-            )
         else:
             raise ValueError(f"unknown backend: {name}")
diff --git a/nnsmith/backends/iree.py b/nnsmith/backends/iree.py
deleted file mode 100644
index fcecc6f..0000000
--- a/nnsmith/backends/iree.py
+++ /dev/null
@@ -1,71 +0,0 @@
-from typing import Dict
-
-import iree.compiler.tf
-import iree.runtime
-import numpy as np
-import tensorflow as tf
-from importlib_metadata import version
-from multipledispatch import dispatch
-
-from nnsmith.backends.factory import BackendCallable, BackendFactory
-from nnsmith.materialize.tensorflow import TFModel
-
-
-class IREEFactory(BackendFactory):
-    def __init__(self, target="cpu", optmax: bool = False):
-        """
-        Initialize the IREE backend factory.
-
-        Parameters
-        ----------
-        target : str, optional
-            The compilation target including "cpu" (same as "llvm-cpu"), "vmvx", "vulkan-spirv", by default "cpu"
-        optmax : bool, optional
-            Release mode or not, by default False
-        """
-        if target == "cpu":
-            target = "llvm-cpu"
-        supported_backends = [
-            "vmvx",
-            "llvm-cpu",
-            "vulkan-spirv",
-        ]
-        assert (
-            target in supported_backends
-        ), f"Unsupported target {target}. Consider one of {supported_backends}"
-        super().__init__(target=target, optmax=optmax)
-
-    @property
-    def system_name(self) -> str:
-        return "iree"
-
-    @property
-    def version(self) -> str:
-        return version("iree-compiler")
-
-    @dispatch(TFModel)
-    def make_backend(self, model: TFModel) -> BackendCallable:
-        setattr(
-            model.net,
-            "iree_fn",
-            tf.function(input_signature=[model.input_specs])(model.net.call_by_dict),
-        )
-        # https://iree-python-api.readthedocs.io/en/latest/compiler/tools.html
-        vm_flatbuffer = iree.compiler.tf.compile_module(
-            model.net,
-            target_backends=[self.target],
-            exported_names=["iree_fn"],
-            optimize=self.optmax,
-        )
-
-        compiled_model = iree.runtime.load_vm_flatbuffer(
-            vm_flatbuffer,
-            backend=self.target,
-        )
-
-        def closure(inputs: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
-            outputs = compiled_model.iree_fn(inputs)  # can directly accept np.ndarray
-            outputs = {k: np.array(v) for k, v in outputs.items()}
-            return outputs
-
-        return closure
diff --git a/nnsmith/backends/tflite.py b/nnsmith/backends/tflite.py
index 2066c7c..c601d60 100644
--- a/nnsmith/backends/tflite.py
+++ b/nnsmith/backends/tflite.py
@@ -1,10 +1,11 @@
 from os import PathLike
-from typing import Dict
+from typing import Dict, List
 
 import numpy as np
 import tensorflow as tf  # type: ignore
 from multipledispatch import dispatch
 
+from nnsmith.abstract.dtype import DType
 from nnsmith.backends.factory import BackendCallable, BackendFactory
 from nnsmith.materialize.tensorflow import TFModel, TFNetCallable
 
@@ -154,3 +155,9 @@ def load_content(self, path: PathLike) -> bytes:
 
     def load_backend(self, path: PathLike) -> BackendCallable:
         return self.make_backend_from_content(self.load_content(path))
+
+    @classmethod
+    def skip_dtypes(cls) -> List[DType]:
+        # PyTorch-ONNX current does not support fully complex (e.g., scalar).
+        # uint8 crash over Equal (revert after their fix).
+        return [DType.uint8]
diff --git a/nnsmith/graph_gen.py b/nnsmith/graph_gen.py
index dddf603..13d519b 100644
--- a/nnsmith/graph_gen.py
+++ b/nnsmith/graph_gen.py
@@ -58,14 +58,21 @@ def __init__(
         self.concr_ph_dim_rng = concr_ph_dim_rng
         self.max_elem_per_tensor = max_elem_per_tensor
         self.rank_choices = rank_choices if rank_choices else rank_all()
+        # analyze the dtypes used by the opset
+        dtype_top = set()
+        for op in opset:
+            dtype_top.update({dt for dtc in op.in_dtypes + op.out_dtypes for dt in dtc})
+
         self.dtype_choices = (
             [
                 dt if isinstance(dt, DType) else DType.from_str(dt)
                 for dt in dtype_choices
             ]
             if dtype_choices
-            else DTYPE_ALL
+            else DTYPE_GEN_ALL
         )
+
+        self.dtype_choices = list(dtype_top.intersection(self.dtype_choices))
         assert len(self.dtype_choices) > 0, "dtype_choices must not be empty"
         assert len(self.rank_choices) > 0, "rank_choices must not be empty"
 
@@ -105,8 +112,7 @@ def make_random_concrete_placeholder(self, rank, dtype=None):
             shape.append(v)
             product *= v
 
-        # shuffle
-        random.shuffle(shape)
+        random.shuffle(shape)  # shuffle
 
         ph = Placeholder(
             AbsTensor(
@@ -127,10 +133,10 @@ def random_dtype_gen(self):
         wts = [1] * len(dtypes)
         for dt in DTYPE_GEN_FLOATS:
             if dt in dtypes:
-                wts[DTYPE_GEN_ALL.index(dt)] = 4
+                wts[dtypes.index(dt)] = 4
         for dt in DTYPE_GEN_INTS:
             if dt in dtypes:
-                wts[DTYPE_GEN_ALL.index(dt)] = 1
+                wts[dtypes.index(dt)] = 1
         return random.choices(dtypes, weights=wts)[0]
 
     def new_sym(self, name):
diff --git a/nnsmith/materialize/__init__.py b/nnsmith/materialize/__init__.py
index bd8303f..0efd589 100644
--- a/nnsmith/materialize/__init__.py
+++ b/nnsmith/materialize/__init__.py
@@ -9,6 +9,7 @@
 import numpy as np
 from multipledispatch import dispatch
 
+from nnsmith.abstract.dtype import DType
 from nnsmith.abstract.op import AbsOpBase, Constant
 from nnsmith.abstract.tensor import AbsTensor
 from nnsmith.error import SanityCheck
@@ -172,6 +173,10 @@ def attach_viz(self, ir: GraphIR) -> None:
     def dump_viz(self, path: PathLike) -> None:
         viz_dot(self.dotstring, path)
 
+    @staticmethod
+    def skip_dtypes() -> List[DType]:
+        return []
+
     @staticmethod
     def init(name, backend_target=None) -> Type["Model"]:
         if name is None:
diff --git a/nnsmith/materialize/onnx/__init__.py b/nnsmith/materialize/onnx/__init__.py
index e8ae68d..ae9eab5 100644
--- a/nnsmith/materialize/onnx/__init__.py
+++ b/nnsmith/materialize/onnx/__init__.py
@@ -3,7 +3,7 @@
 import warnings
 from io import BytesIO
 from os import PathLike
-from typing import Dict, Optional, Tuple, Type, Union
+from typing import Dict, List, Optional, Tuple, Type, Union
 
 import onnx
 import onnx.checker
@@ -13,7 +13,7 @@
 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.dtype import DTYPE_GEN_COMPLEX, DType
 from nnsmith.abstract.op import AbsTensor
 from nnsmith.gir import GraphIR
 from nnsmith.macro import onnx2external_data_dir
@@ -274,3 +274,7 @@ def get_onnx_from_torch(self) -> onnx.ModelProto:
     @staticmethod
     def name_suffix() -> str:
         return ".onnx"
+
+    @classmethod
+    def skip_dtypes(cls) -> List[DType]:
+        return DTYPE_GEN_COMPLEX
diff --git a/nnsmith/materialize/tensorflow/dialect.py b/nnsmith/materialize/tensorflow/dialect.py
index 6c8931b..84ed7ec 100644
--- a/nnsmith/materialize/tensorflow/dialect.py
+++ b/nnsmith/materialize/tensorflow/dialect.py
@@ -4,8 +4,6 @@
 from nnsmith.abstract.arith import *
 from nnsmith.abstract.dtype import DType
 from nnsmith.abstract.op import (
-    FLOPS_LIM,
-    Z3_CONS_FLOPS,
     BcastBinaryOp,
     ElementWiseUnaryOp,
     MatMul,
@@ -180,9 +178,6 @@ def requires(self, input_shapes):
                 ),
             )
         )
-        # limit FLOPS
-        if Z3_CONS_FLOPS:
-            cons.append(nnsmith_le(self.flops(input_shapes), FLOPS_LIM))
         return cons
 
     def flops(self, input_shapes):
diff --git a/nnsmith/materialize/tensorflow/forward.py b/nnsmith/materialize/tensorflow/forward.py
index 0293a42..cd2511c 100644
--- a/nnsmith/materialize/tensorflow/forward.py
+++ b/nnsmith/materialize/tensorflow/forward.py
@@ -195,24 +195,20 @@ def forward_fn(op: Softmax):
 @operator_impl(Slice)
 def forward_fn(op: Slice):
     reg = op.extra_attrs["region"]
+    shape = op.input_like[0].shape
+    dim_s = shape[op.extra_attrs["axis"]]
+    start, end = op.start, op.end
+    if reg in ["left", "mid"]:
+        start -= dim_s
+    # actual end would be 0, which is not really 'left'
+    if reg == "left" and end < dim_s and end != Slice.INT_MAX:
+        end -= dim_s
+    idx = tuple(
+        slice(None, None) if i != op.extra_attrs["axis"] else slice(start, end, op.step)
+        for i in range(op.extra_attrs["ndims"])
+    )
 
-    def _slice(x):
-        dim_s = x.shape[op.extra_attrs["axis"]]
-        start, end = op.start, op.end
-        if reg in ["left", "mid"]:
-            start -= dim_s
-        # actual end would be 0, which is not really 'left'
-        if reg == "left" and end < dim_s and end != Slice.INT_MAX:
-            end -= dim_s
-        s = tuple(
-            slice(None, None)
-            if i != op.extra_attrs["axis"]
-            else slice(start, end, op.step)
-            for i in range(op.extra_attrs["ndims"])
-        )
-        return x[s]
-
-    return _slice
+    return lambda x: x[idx]
 
 
 @operator_impl(BatchNorm2d)
@@ -235,8 +231,9 @@ def _reshape(x):
 @operator_impl(Transpose)
 def forward_fn(op: Transpose):
     def _transpose(x: tf.Tensor):
-        dim0, dim1 = op._init_swap_dims(list(x.shape))
-        perm = list(range(len(x.shape)))
+        aten = op.input_like[0]
+        dim0, dim1 = op._init_swap_dims(aten.shape)
+        perm = list(range(aten.ndims))
         perm[dim0], perm[dim1] = perm[dim1], perm[dim0]
         return tf.transpose(x, perm=perm)
 
diff --git a/nnsmith/materialize/torch/dialect.py b/nnsmith/materialize/torch/dialect.py
index e34f00b..c0dff75 100644
--- a/nnsmith/materialize/torch/dialect.py
+++ b/nnsmith/materialize/torch/dialect.py
@@ -3,7 +3,14 @@
 from typing import List, Tuple, Union
 
 from nnsmith.abstract.arith import *
-from nnsmith.abstract.dtype import DTYPE_GEN_ALL, DTYPE_GEN_NON_BOOL, DType
+from nnsmith.abstract.dtype import (
+    DTYPE_GEN_ALL,
+    DTYPE_GEN_COMPLEX,
+    DTYPE_GEN_FLOATS,
+    DTYPE_GEN_INTS,
+    DTYPE_GEN_NON_BOOL,
+    DType,
+)
 from nnsmith.abstract.op import ReduceBase, UnaryOpBase, mark_materialize, rank_from
 from nnsmith.abstract.tensor import AbsTensor
 from nnsmith.error import ConstraintCheck
@@ -76,11 +83,10 @@ def deduct_inp_ranks_and_dtype(
 @mark_materialize("torch")
 class TorchReduceSum(ReduceBase):
     in_dtypes = [(i,) for i in DTYPE_GEN_NON_BOOL]
-    out_dtypes = [(i,) for i in DTYPE_GEN_NON_BOOL]
+    out_dtypes = [(i,) for i in DTYPE_GEN_FLOATS + DTYPE_GEN_COMPLEX + [DType.int64]]
 
     def type_transfer(self, input_shapes: List[AbsTensor]) -> List[AbsTensor]:
         output = super().type_transfer(input_shapes)
-        # This is a PyTorch trick...
-        if input_shapes[0].dtype == DType.int32:
+        if input_shapes[0].dtype in DTYPE_GEN_INTS:  # This is a PyTorch trick...
             output[0].dtype = DType.int64
         return output
diff --git a/nnsmith/materialize/torch/forward.py b/nnsmith/materialize/torch/forward.py
index e59cfe5..e166f4c 100644
--- a/nnsmith/materialize/torch/forward.py
+++ b/nnsmith/materialize/torch/forward.py
@@ -22,7 +22,7 @@
 @operator_impl(Constant)
 def forward_fn(op: Constant):
     class ConstFn(torch.nn.Module):
-        def __init__(self, data) -> None:
+        def __init__(self, data: torch.Tensor) -> None:
             super().__init__()
             self.data = torch.nn.parameter.Parameter(
                 data, requires_grad=data.is_floating_point()
@@ -36,17 +36,17 @@ def forward(self):
 
 @operator_impl(ReLU)
 def forward_fn(op: ReLU):
-    return torch.nn.ReLU()
+    return torch.relu
 
 
 @operator_impl(GELU)
 def forward_fn(op: GELU):
-    return torch.nn.GELU()
+    return torch.nn.functional.gelu
 
 
 @operator_impl(LeakyReLU)
 def forward_fn(op: LeakyReLU):
-    return torch.nn.LeakyReLU(op.negative_slope)
+    return torch.nn.LeakyReLU(0.01)
 
 
 @operator_impl(PReLU)
@@ -56,7 +56,7 @@ def forward_fn(op: PReLU):
 
 @operator_impl(Sigmoid)
 def forward_fn(op: Sigmoid):
-    return torch.nn.Sigmoid()
+    return torch.nn.functional.sigmoid
 
 
 @operator_impl(Sin)
@@ -117,11 +117,9 @@ def forward_fn(op: Mul):
 
 @operator_impl(Div)
 def forward_fn(op: Div):
-    return lambda up, down: torch.div(
-        up,
-        down,
-        rounding_mode="floor" if DType.from_torch(up.dtype) in DTYPE_GEN_INTS else None,
-    )
+    if op.input_like[0].dtype in DTYPE_GEN_INTS:
+        return lambda up, down: torch.div(up, down, rounding_mode="floor")
+    return torch.div
 
 
 @operator_impl(Max)
@@ -235,24 +233,20 @@ def forward_fn(op: AvgPool2d):
 @operator_impl(Slice)
 def forward_fn(op: Slice):
     reg = op.extra_attrs["region"]
+    shape = op.input_like[0].shape
+    dim_s = shape[op.extra_attrs["axis"]]
+    start, end = op.start, op.end
+    if reg in ["left", "mid"]:
+        start -= dim_s
+    # actual end would be 0, which is not really 'left'
+    if reg == "left" and end < dim_s and end != Slice.INT_MAX:
+        end -= dim_s
+    idx = tuple(
+        slice(None, None) if i != op.extra_attrs["axis"] else slice(start, end, op.step)
+        for i in range(op.extra_attrs["ndims"])
+    )
 
-    def _func(x):
-        dim_s = x.shape[op.extra_attrs["axis"]]
-        start, end = op.start, op.end
-        if reg in ["left", "mid"]:
-            start -= dim_s
-        # actual end would be 0, which is not really 'left'
-        if reg == "left" and end < dim_s and end != Slice.INT_MAX:
-            end -= dim_s
-        s = tuple(
-            slice(None, None)
-            if i != op.extra_attrs["axis"]
-            else slice(start, end, op.step)
-            for i in range(op.extra_attrs["ndims"])
-        )
-        return x[s]
-
-    return _func
+    return lambda x: x[idx]
 
 
 @operator_impl(Pad)
@@ -269,13 +263,9 @@ def forward_fn(op: Pad):
         )
 
 
-def type_from_torch_tensor(tensor: torch.Tensor):
-    return AbsTensor(list(tensor.shape), DType.from_torch(tensor.dtype))
-
-
 @operator_impl(Expand)
 def forward_fn(op: Expand):
-    return lambda x: x.expand(*op.type_transfer([type_from_torch_tensor(x)])[0].shape)
+    return lambda x: x.expand(*op.output_like[0].shape)
 
 
 @operator_impl(BatchNorm2d)
@@ -314,16 +304,13 @@ def forward_fn(op: Reshape):
 
 @operator_impl(Flatten)
 def forward_fn(op: Flatten):
-    return lambda x: x.flatten()
+    return torch.Tensor.flatten
 
 
 @operator_impl(Transpose)
 def forward_fn(op: Transpose):
-    def f(x: torch.Tensor):
-        dim0, dim1 = op._init_swap_dims(list(x.shape))
-        return x.transpose(dim0, dim1)
-
-    return f
+    dim0, dim1 = op._init_swap_dims(op.input_like[0].shape)
+    return lambda x: x.transpose(dim0, dim1)
 
 
 # NearestInterp
@@ -359,15 +346,14 @@ def forward_fn(op: TrilinearInterp):
 def forward_fn(op: Squeeze):
     if op.extra_attrs["reduce_dim"] is not None:
         return lambda x: x.squeeze(op.extra_attrs["reduce_dim"])
-    else:
-        return lambda x: x.squeeze()
+    return torch.Tensor.squeeze
 
 
 @operator_impl(TorchReduceSum)
 def forward_fn(op: TorchReduceSum):
     if op.extra_attrs["reduce_dim"] is not None:
         return lambda x: x.sum(op.extra_attrs["reduce_dim"])
-    return lambda x: x.sum()
+    return torch.Tensor.sum
 
 
 # ReduceMin
@@ -375,7 +361,7 @@ def forward_fn(op: TorchReduceSum):
 def forward_fn(op: ReduceMin):
     if op.extra_attrs["reduce_dim"] is not None:
         return lambda x: x.min(op.extra_attrs["reduce_dim"]).values
-    return lambda x: x.min()
+    return torch.Tensor.min
 
 
 # ReduceMax
@@ -383,7 +369,7 @@ def forward_fn(op: ReduceMin):
 def forward_fn(op: ReduceMax):
     if op.extra_attrs["reduce_dim"] is not None:
         return lambda x: x.max(op.extra_attrs["reduce_dim"]).values
-    return lambda x: x.max()
+    return torch.Tensor.max
 
 
 # ReduceMean
@@ -391,7 +377,7 @@ def forward_fn(op: ReduceMax):
 def forward_fn(op: ReduceMean):
     if op.extra_attrs["reduce_dim"] is not None:
         return lambda x: x.mean(op.extra_attrs["reduce_dim"])
-    return lambda x: x.mean()
+    return torch.Tensor.mean
 
 
 # ArgMin
@@ -399,7 +385,7 @@ def forward_fn(op: ReduceMean):
 def forward_fn(op: ArgMin):
     if op.extra_attrs["reduce_dim"] is not None:
         return lambda x: x.argmin(op.extra_attrs["reduce_dim"])
-    return lambda x: x.argmin()
+    return torch.Tensor.argmin
 
 
 # ArgMax
@@ -407,7 +393,7 @@ def forward_fn(op: ArgMin):
 def forward_fn(op: ArgMax):
     if op.extra_attrs["reduce_dim"] is not None:
         return lambda x: x.argmax(op.extra_attrs["reduce_dim"])
-    return lambda x: x.argmax()
+    return torch.Tensor.argmax
 
 
 # Tril
diff --git a/nnsmith/materialize/torch/symbolnet.py b/nnsmith/materialize/torch/symbolnet.py
index f436d06..4b9fe24 100644
--- a/nnsmith/materialize/torch/symbolnet.py
+++ b/nnsmith/materialize/torch/symbolnet.py
@@ -7,11 +7,13 @@
 import torch
 from torch import nn
 
+from nnsmith.abstract.dtype import DType
 from nnsmith.abstract.op import AbsOpBase, Input
+from nnsmith.abstract.tensor import AbsTensor
 from nnsmith.error import ConstraintCheck, ConstraintError, SanityCheck
 from nnsmith.gir import GraphIR
 from nnsmith.logging import TORCH_LOG
-from nnsmith.materialize.torch.forward import forward_fn, type_from_torch_tensor
+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
 
@@ -31,10 +33,10 @@ def check_type(op: AbsOpBase, tensors, is_input=True, msg=""):
         ), f"{op}'s {ioro} has {len(like)} abs. input, but got {len(tensors)} real inputs."
 
         for i, ten in enumerate(tensors):
-            ttype = type_from_torch_tensor(ten)
+            ttype = AbsTensor(list(ten.shape), DType.from_torch(ten.dtype))
             assert (
                 like[i] == ttype
-            ), f"{msg} {ioro} abstract type != concrete {type_from_torch_tensor(ten)} for {op} {op.input_like} -> {op.output_like}"
+            ), f"{msg} {ioro} abstract type != concrete {ttype} for {op} {op.input_like} -> {op.output_like}"
 
 
 # Probablistically, sampling at positive domain is beneficial.
diff --git a/nnsmith/narrow_spec.py b/nnsmith/narrow_spec.py
index 9481a55..0c6bea4 100644
--- a/nnsmith/narrow_spec.py
+++ b/nnsmith/narrow_spec.py
@@ -116,7 +116,7 @@ def infer_topset_from_scratch(
         for i, ranks in enumerate(op.inp_ranks):
             if op.same_inp_dims and inputs:
                 rank = inputs[0].ndims
-            else:
+            else:  # FIXME(@ganler): consider rank check over scalar & non-scalar.
                 rank = min(ranks)
             shape = AbsTensor(
                 shape=[z3.Int("s%s" % (k)) for k in range(rank)],
@@ -145,6 +145,13 @@ def infer_topset_from_scratch(
             concrete_op, concre_input_shapes, available_idtypes
         )
 
+        # filter out unsupported dtypes by model.
+        single_op_irs = [
+            sset
+            for sset in single_op_irs
+            if set(model_cls.skip_dtypes()).isdisjoint(sset[0] + sset[1])
+        ]
+
         if factory:
             single_op_irs = [
                 sset
@@ -175,6 +182,7 @@ def infer_topset_from_scratch(
             else:  # Test model dumping
                 with tempfile.TemporaryDirectory() as tmpdirname:
                     try:
+                        model.make_oracle()  # try-run.
                         model_path = os.path.join(
                             tmpdirname, model.name_prefix() + model.name_suffix()
                         )
diff --git a/requirements/sys/iree.txt b/requirements/sys/iree.txt
deleted file mode 100644
index 2e686a7..0000000
--- a/requirements/sys/iree.txt
+++ /dev/null
@@ -1,5 +0,0 @@
--f https://github.com/iree-org/iree/releases
-iree-compiler
-iree-runtime
-iree-tools-tf
-importlib_metadata
diff --git a/setup.cfg b/setup.cfg
index 171c742..69c2a55 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -31,10 +31,6 @@ torch = torch
 tvm = apache-tvm
       torch
       onnx
-iree = tf-nightly
-       iree-compiler
-       iree-runtime
-       iree-tools-tf
 
 [options.package_data]
 nnsmith = config/**/*.yaml


Test patch:
diff --git a/tests/tensorflow/test_iree_backend.py b/tests/tensorflow/test_iree_backend.py
deleted file mode 100644
index fa85f30..0000000
--- a/tests/tensorflow/test_iree_backend.py
+++ /dev/null
@@ -1,53 +0,0 @@
-import pytest
-
-from nnsmith.abstract.dtype import DType
-from nnsmith.backends import BackendFactory
-from nnsmith.graph_gen import model_gen
-from nnsmith.materialize import Model, TestCase
-from nnsmith.narrow_spec import auto_opconfig, auto_opset
-
-TestCase.__test__ = False  # supress PyTest warning
-
-
-def test_narrow_spec_cache_make_and_reload():
-    factory = BackendFactory.init("iree", target="cpu", optmax=True)
-    ModelType = Model.init("tensorflow")
-    opset_lhs = auto_opconfig(ModelType, factory)
-    assert opset_lhs, "Should not be empty... Something must go wrong."
-    opset_rhs = auto_opconfig(ModelType, factory)
-    assert opset_lhs == opset_rhs
-
-    # Assert types
-    assert isinstance(opset_lhs["core.ReLU"].in_dtypes[0][0], DType)
-
-    # Assert Dictionary Type Equality
-    assert type(opset_lhs) == type(opset_rhs)
-    assert type(opset_lhs["core.ReLU"]) == type(opset_rhs["core.ReLU"])
-    assert type(opset_lhs["core.ReLU"].in_dtypes[0][0]) == type(
-        opset_rhs["core.ReLU"].in_dtypes[0][0]
-    )
-
-
-def test_synthesized_tf_model(tmp_path):
-    d = tmp_path / "test_iree"
-    d.mkdir()
-
-    targets = ["cpu"]
-
-    ModelType = Model.init("tensorflow")
-    for target in targets:
-        factory = BackendFactory.init("iree", target=target, optmax=False)
-
-        gen = model_gen(
-            opset=auto_opset(ModelType, factory),
-            seed=23132,
-            max_nodes=4,
-        )  # One op should not be easily wrong... I guess.
-
-        model = ModelType.from_gir(gen.make_concrete())
-
-        oracle = model.make_oracle()
-
-        testcase = TestCase(model, oracle)
-        assert factory.verify_testcase(testcase) is None
-        testcase.dump(root_folder=d)