diff --git a/cinn/hlir/op/contrib/sort.cc b/cinn/hlir/op/contrib/sort.cc
index e8471bcb4a..0f324a73f2 100644
--- a/cinn/hlir/op/contrib/sort.cc
+++ b/cinn/hlir/op/contrib/sort.cc
@@ -364,7 +364,7 @@ std::vector<std::vector<int>> InferShapeForTopK(const std::vector<std::vector<in
   }
   CHECK_GE(axis, 0);
   CHECK_LT(axis, res[0].size());
-  res[0][axis] = k;
+  res[0][axis] = std::min(res[0][axis], k);
   return {res[0], res[0]};
 }
 
diff --git a/python/tests/ops/test_top_k_op.py b/python/tests/ops/test_top_k_op.py
index b557797117..0241f02cb2 100644
--- a/python/tests/ops/test_top_k_op.py
+++ b/python/tests/ops/test_top_k_op.py
@@ -14,76 +14,294 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import unittest
-import numpy as np
-from op_test import OpTest, OpTestTool
 import paddle
-import paddle.nn.functional as F
-import cinn
 from cinn.frontend import *
 from cinn.common import *
+from op_test import OpTest, OpTestTool
+from op_test_helper import TestCaseHelper, run_test
 
 
 @OpTestTool.skip_if(not is_compiled_with_cuda(),
                     "x86 test will be skipped due to timeout.")
 class TestTopKOp(OpTest):
     def setUp(self):
-        self.init_case()
+        print(f"\nRunning {self.__class__.__name__}: {self.case}")
+        self.inputs = {}
+        self.prepare_inputs()
 
-    def init_case(self):
+    def prepare_inputs(self):
         self.inputs = {
-            "x1": np.random.random([
-                2,
-                4,
-            ]).astype("float32")
+            "x": self.random(self.case["shape"], self.case["dtype"])
         }
-        self.k = 1
-        self.axis = 1
+        self.k = self.case["k"]
+        self.axis = self.case["axis"]
+        self.largest = self.case["largest"]
 
     def build_paddle_program(self, target):
-        axis = -1
-        x1 = paddle.to_tensor(self.inputs["x1"], stop_gradient=True)
-        out = paddle.topk(x1, self.k, self.axis)
-
+        x = paddle.to_tensor(self.inputs["x"], stop_gradient=True)
+        if x.shape[self.axis] < self.k:
+            self.k = x.shape[self.axis]
+        out = paddle.topk(x, self.k, self.axis)
         self.paddle_outputs = [out[0], out[1]]
 
     def build_cinn_program(self, target):
-        builder = NetBuilder("sum")
-        x1 = builder.create_input(Float(32), self.inputs["x1"].shape, "x1")
-        out = builder.top_k(x1, self.k, self.axis, False)
+        builder = NetBuilder("topk")
+        x = builder.create_input(
+            self.nptype2cinntype(self.inputs["x"].dtype),
+            self.inputs["x"].shape, "x")
+        out = builder.top_k(x, self.k, self.axis, self.largest)
         prog = builder.build()
         forward_res = self.get_cinn_output(
-            prog, target, [x1], [self.inputs["x1"]], [out[0], out[1]])
-
+            prog, target, [x], [self.inputs["x"]], [out[0], out[1]])
         self.cinn_outputs = forward_res
 
     def test_check_results(self):
         self.check_outputs_and_grads()
 
 
-class TestTopKCase1(TestTopKOp):
-    def init_case(self):
-        self.inputs = {
-            "x1": np.random.random([
-                2,
-                4,
-            ]).astype("float32")
-        }
-        self.k = 2
-        self.axis = 1
+class TestTopKOpDumpicateElement(TestTopKOp):
+    def setUp(self):
+        self.inputs = {}
+        self.prepare_inputs()
 
+    def prepare_inputs(self):
+        self.inputs = {"x": self.random([128], "int64", -10, 10)}
+        self.axis = 0
+        self.largest = False
+        self.k = 5
 
-class TestTopKCase2(TestTopKOp):
-    def init_case(self):
-        self.inputs = {
-            "x1": np.random.random([
-                2,
-                4,
-            ]).astype("float32")
-        }
-        self.k = 1
+
+# known issue: same as sort op
+# This test case will cause CINN to allocate a large amount of GPU memory, nearly 10 GB.
+class TestTopKOpLargeCudaMemoryOccupation(TestTopKOp):
+    def setUp(self):
+        self.inputs = {}
+        self.prepare_inputs()
+
+    def prepare_inputs(self):
+        self.inputs = {"x": self.random([8192], "float64")}
         self.axis = 0
+        self.largest = False
+        self.k = 5
+
+
+class TestTopKOpShapeTest(TestCaseHelper):
+    def init_attrs(self):
+        self.class_name = "TestTopKOpShapeTest"
+        self.cls = TestTopKOp
+        self.inputs = [
+            {
+                "shape": [512],
+                "k": 3
+            },
+            {
+                "shape": [1024],
+                "k": 10
+            },
+            {
+                "shape": [1200],
+                "k": 1024
+            },
+            {
+                "shape": [64, 16],
+                "k": 3
+            },
+            {
+                "shape": [4, 32, 8],
+                "k": 4
+            },
+            {
+                "shape": [16, 8, 4, 2],
+                "k": 5
+            },
+            {
+                "shape": [2, 8, 4, 2, 5],
+                "k": 1
+            },
+            {
+                "shape": [4, 8, 1, 2, 16],
+                "k": 3
+            },
+            {
+                "shape": [1],
+                "k": 1
+            },
+            {
+                "shape": [1, 1, 1, 1],
+                "k": 1
+            },
+            {
+                "shape": [1, 1, 1, 1, 1],
+                "k": 1
+            },
+        ]
+        self.dtypes = [{"dtype": "float32"}]
+        self.attrs = [{"axis": 0, "largest": True}]
+
+
+class TestTopKOpDtypeTest(TestCaseHelper):
+    def init_attrs(self):
+        self.class_name = "TestTopKOpDtypeTest"
+        self.cls = TestTopKOp
+        self.inputs = [
+            {
+                "shape": [1024],
+            },
+            {
+                "shape": [64, 16],
+            },
+            {
+                "shape": [4, 32, 8],
+            },
+            {
+                "shape": [16, 8, 4, 2],
+            },
+        ]
+        self.dtypes = [
+            {
+                "dtype": "float32"
+            },
+            {
+                "dtype": "float64"
+            },
+            {
+                "dtype": "int32"
+            },
+            {
+                "dtype": "int64"
+            },
+        ]
+        self.attrs = [{"axis": 0, "largest": True, "k": 3}]
+
+
+class TestTopKOpAxisTest(TestCaseHelper):
+    def init_attrs(self):
+        self.class_name = "TestTopKOpAxisTest"
+        self.cls = TestTopKOp
+        self.inputs = [
+            {
+                "shape": [16, 8, 4, 8],
+            },
+        ]
+        self.dtypes = [{"dtype": "float32"}]
+        self.attrs = [
+            {
+                "axis": 0,
+                "largest": True,
+                "k": 3
+            },
+            {
+                "axis": 1,
+                "largest": True,
+                "k": 3
+            },
+            {
+                "axis": 2,
+                "largest": True,
+                "k": 3
+            },
+            {
+                "axis": 3,
+                "largest": True,
+                "k": 3
+            },
+        ]
+
+
+class TestTopKOpKTest(TestCaseHelper):
+    def init_attrs(self):
+        self.class_name = "TestTopKOpKTest"
+        self.cls = TestTopKOp
+        self.inputs = [
+            {
+                "shape": [16, 8, 4, 2],
+            },
+        ]
+        self.dtypes = [{"dtype": "float32"}]
+        self.attrs = [
+            {
+                "axis": 0,
+                "largest": True,
+                "k": 8
+            },
+            {
+                "axis": 1,
+                "largest": True,
+                "k": 4
+            },
+            {
+                "axis": 2,
+                "largest": True,
+                "k": 2
+            },
+            {
+                "axis": 3,
+                "largest": True,
+                "k": 1
+            },
+            {
+                "axis": 0,
+                "largest": True,
+                "k": 20
+            },
+            {
+                "axis": 1,
+                "largest": True,
+                "k": 10
+            },
+            {
+                "axis": 2,
+                "largest": True,
+                "k": 10
+            },
+            {
+                "axis": 3,
+                "largest": True,
+                "k": 5
+            },
+        ]
+
+
+class TestTopKOpAscendingTest(TestTopKOpShapeTest):
+    def init_attrs(self):
+        self.class_name = "TestTopKOpAscendingTest"
+        self.cls = TestTopKOp
+        self.inputs = [
+            {
+                "shape": [16, 8, 4, 8],
+            },
+        ]
+        self.dtypes = [{"dtype": "float32"}]
+        self.attrs = [
+            {
+                "axis": 0,
+                "largest": False,
+                "k": 3
+            },
+            {
+                "axis": 1,
+                "largest": False,
+                "k": 3
+            },
+            {
+                "axis": 2,
+                "largest": False,
+                "k": 3
+            },
+            {
+                "axis": 3,
+                "largest": False,
+                "k": 3
+            },
+        ]
 
 
 if __name__ == "__main__":
-    unittest.main()
+    run_test(TestTopKOpDumpicateElement)
+    # run_test(TestTopKOpLargeCudaMemoryOccupation)
+
+    TestTopKOpShapeTest().run()
+    TestTopKOpDtypeTest().run()
+    TestTopKOpAxisTest().run()
+    TestTopKOpKTest().run()
+    TestTopKOpAscendingTest().run()
diff --git a/python/tests/ops/test_transpose_op.py b/python/tests/ops/test_transpose_op.py
index 7153a2dcf1..13e6564a3a 100644
--- a/python/tests/ops/test_transpose_op.py
+++ b/python/tests/ops/test_transpose_op.py
@@ -14,25 +14,26 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import cinn
-import numpy as np
 import paddle
-import unittest
-
 from cinn.frontend import *
 from cinn.common import *
 from op_test import OpTest, OpTestTool
+from op_test_helper import TestCaseHelper
 
 
 @OpTestTool.skip_if(not is_compiled_with_cuda(),
                     "x86 test will be skipped due to timeout.")
 class TestTransposeOp(OpTest):
     def setUp(self):
-        self.init_case()
+        print(f"\nRunning {self.__class__.__name__}: {self.case}")
+        self.inputs = {}
+        self.prepare_inputs()
 
-    def init_case(self):
-        self.inputs = {"x": self.random([2, 3], "float32")}
-        self.axes = [1, 0]
+    def prepare_inputs(self):
+        self.inputs = {
+            "x": self.random(self.case["shape"], self.case["dtype"])
+        }
+        self.axes = self.case["axes"]
 
     def build_paddle_program(self, target):
         x = paddle.to_tensor(self.inputs["x"], stop_gradient=True)
@@ -45,24 +46,227 @@ def build_paddle_program(self, target):
     def build_cinn_program(self, target):
         builder = NetBuilder("transpose_test")
         x = builder.create_input(
-            self.nptype2cinntype(self.inputs['x'].dtype),
+            self.nptype2cinntype(self.inputs["x"].dtype),
             self.inputs["x"].shape, "x")
         out = builder.transpose(x, self.axes)
 
         prog = builder.build()
         res = self.get_cinn_output(prog, target, [x], [self.inputs["x"]],
                                    [out])
-        self.cinn_outputs = [res[0]]
+        self.cinn_outputs = res
 
     def test_check_results(self):
         self.check_outputs_and_grads(all_equal=True)
 
 
-class TestTransposeOpWithNegAxis(TestTransposeOp):
-    def init_case(self):
-        self.inputs = {"x": self.random([10, 8, 2], "float32")}
-        self.axes = [-1, 1, -3]
+class TestTransposeOpShapeTest(TestCaseHelper):
+    def init_attrs(self):
+        self.class_name = "TestTransposeOpShapeTest"
+        self.cls = TestTransposeOp
+        self.inputs = [
+            {
+                "shape": [512],
+                "axes": [0]
+            },
+            {
+                "shape": [1024],
+                "axes": [-1]
+            },
+            {
+                "shape": [1200],
+                "axes": [0]
+            },
+            {
+                "shape": [64, 16],
+                "axes": [1, 0]
+            },
+            {
+                "shape": [4, 32, 8],
+                "axes": [1, 0, 2]
+            },
+            {
+                "shape": [16, 8, 4, 2],
+                "axes": [0, 2, 1, 3]
+            },
+            {
+                "shape": [2, 8, 4, 2, 5],
+                "axes": [0, 2, 3, 1, 4]
+            },
+            {
+                "shape": [4, 8, 1, 2, 16],
+                "axes": [0, 2, 4, 1, 3]
+            },
+            {
+                "shape": [1],
+                "axes": [0]
+            },
+            {
+                "shape": [1, 1, 1, 1],
+                "axes": [0, 2, 3, 1]
+            },
+            {
+                "shape": [1, 1, 1, 1, 1],
+                "axes": [0, 2, 3, 1, 4]
+            },
+        ]
+        self.dtypes = [{"dtype": "float32"}]
+        self.attrs = []
+
+
+class TestTransposeOpOnesTest(TestCaseHelper):
+    def init_attrs(self):
+        self.class_name = "TestTransposeOpOnesTest"
+        self.cls = TestTransposeOp
+        self.inputs = [
+            {
+                "shape": [1],
+                "axes": [0]
+            },
+            {
+                "shape": [1, 1, 1, 1],
+                "axes": [0, 2, 3, 1]
+            },
+            {
+                "shape": [1, 1, 1, 1, 1],
+                "axes": [0, 2, 3, 1, 4]
+            },
+            {
+                "shape": [1, 1, 512, 1, 1],
+                "axes": [0, 2, 3, 1, 4]
+            },
+            {
+                "shape": [1, 2048, 1, 1],
+                "axes": [0, 2, 3, 1]
+            },
+        ]
+        self.dtypes = [{"dtype": "float32"}]
+        self.attrs = []
+
+
+class TestTransposeOpLargeTest(TestCaseHelper):
+    def init_attrs(self):
+        self.class_name = "TestTransposeOpLargeTest"
+        self.cls = TestTransposeOp
+        self.inputs = [
+            {
+                "shape": [2048],
+                "axes": [0]
+            },
+            {
+                "shape": [1, 1, 65536, 1],
+                "axes": [0, 2, 3, 1]
+            },
+            {
+                "shape": [1, 1, 131072, 1, 1],
+                "axes": [0, 2, 3, 1, 4]
+            },
+            {
+                "shape": [1, 1048576, 1, 1],
+                "axes": [0, 2, 3, 1]
+            },
+            {
+                "shape": [16, 32, 64, 32],
+                "axes": [0, 2, 3, 1]
+            },
+        ]
+        self.dtypes = [{"dtype": "float32"}]
+        self.attrs = []
+
+
+class TestTransposeOpDtypeTest(TestCaseHelper):
+    def init_attrs(self):
+        self.class_name = "TestTransposeOpDtypeTest"
+        self.cls = TestTransposeOp
+        self.inputs = [
+            {
+                "shape": [1024],
+                "axes": [0]
+            },
+            {
+                "shape": [64, 16],
+                "axes": [1, 0]
+            },
+            {
+                "shape": [4, 32, 8],
+                "axes": [0, 2, 1]
+            },
+            {
+                "shape": [16, 8, 4, 2],
+                "axes": [1, 2, 3, 0]
+            },
+        ]
+        self.dtypes = [
+            {
+                "dtype": "float16"
+            },
+            {
+                "dtype": "float32"
+            },
+            {
+                "dtype": "float64"
+            },
+            {
+                "dtype": "bool"
+            },
+            {
+                "dtype": "int32"
+            },
+            {
+                "dtype": "int64"
+            },
+        ]
+        self.attrs = []
+
+
+class TestTransposeOpAttributeAxes(TestCaseHelper):
+    def init_attrs(self):
+        self.class_name = "TestTransposeOpAttributeAxes"
+        self.cls = TestTransposeOp
+        self.inputs = [
+            {
+                "shape": [1024],
+                "axes": [0]
+            },
+            {
+                "shape": [1024],
+                "axes": [-1]
+            },
+            {
+                "shape": [64, 16],
+                "axes": [1, 0]
+            },
+            {
+                "shape": [64, 16],
+                "axes": [0, -1]
+            },
+            {
+                "shape": [4, 32, 8],
+                "axes": [0, 2, 1]
+            },
+            {
+                "shape": [4, 32, 8],
+                "axes": [-3, -1, 1]
+            },
+            {
+                "shape": [16, 8, 4, 2],
+                "axes": [1, 2, 3, 0]
+            },
+            {
+                "shape": [16, 8, 4, 2],
+                "axes": [1, -2, -1, -4]
+            },
+        ]
+        self.dtypes = [
+            {
+                "dtype": "float32"
+            },
+        ]
+        self.attrs = []
 
 
 if __name__ == "__main__":
-    unittest.main()
+    TestTransposeOpShapeTest().run()
+    TestTransposeOpOnesTest().run()
+    TestTransposeOpLargeTest().run()
+    TestTransposeOpDtypeTest().run()
+    TestTransposeOpAttributeAxes().run()
diff --git a/python/tests/ops/test_triangular_solve_op.py b/python/tests/ops/test_triangular_solve_op.py
index 616d1fc9a9..f406ca41db 100644
--- a/python/tests/ops/test_triangular_solve_op.py
+++ b/python/tests/ops/test_triangular_solve_op.py
@@ -14,29 +14,30 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import unittest
-import numpy as np
-from op_test import OpTest, OpTestTool
 import paddle
 from cinn.frontend import *
 from cinn.common import *
+from op_test import OpTest, OpTestTool
+from op_test_helper import TestCaseHelper, run_test
 
 
 @OpTestTool.skip_if(not is_compiled_with_cuda(),
                     "triangular solve op support GPU only now.")
 class TestTriangularSolveOp(OpTest):
     def setUp(self):
-        self.init_case()
+        print(f"\nRunning {self.__class__.__name__}: {self.case}")
+        self.inputs = {}
+        self.prepare_inputs()
 
-    def init_case(self):
+    def prepare_inputs(self):
         self.inputs = {
-            "input1": np.random.random((1, 3, 3)).astype(np.float32),
-            "input2": np.random.random((1, 3, 1)).astype(np.float32),
+            "input1": self.random(self.case["shape1"], self.case["dtype"]),
+            "input2": self.random(self.case["shape2"], self.case["dtype"])
         }
-        self.left_side = True
-        self.upper = True
-        self.transpose_a = False
-        self.unit_diagonal = False
+        self.left_side = self.case["left_side"]
+        self.upper = self.case["upper"]
+        self.transpose_a = self.case["transpose_a"]
+        self.unit_diagonal = self.case["unit_diagonal"]
 
     def build_paddle_program(self, target):
         def transpose_last_two_dims(x):
@@ -100,113 +101,207 @@ def test_check_results(self):
         self.check_outputs_and_grads()
 
 
-@OpTestTool.skip_if(not is_compiled_with_cuda(),
-                    "triangular solve op support GPU only now.")
-class TestTriangularSolveOpUnitDiagonal(TestTriangularSolveOp):
-    def init_case(self):
-        self.inputs = {
-            "input1": np.random.random((1, 3, 3)).astype(np.float32),
-            "input2": np.random.random((1, 3, 1)).astype(np.float32),
-        }
-        self.left_side = True
-        self.upper = True
-        self.transpose_a = False
-        self.unit_diagonal = True
-
-
-@OpTestTool.skip_if(not is_compiled_with_cuda(),
-                    "triangular solve op support GPU only now.")
-class TestTriangularSolveOpLower(TestTriangularSolveOp):
-    def init_case(self):
-        self.inputs = {
-            "input1": np.random.random((1, 3, 3)).astype(np.float32),
-            "input2": np.random.random((1, 3, 1)).astype(np.float32),
-        }
-        self.left_side = True
-        self.upper = False
-        self.transpose_a = False
-        self.unit_diagonal = False
-
-
-@OpTestTool.skip_if(not is_compiled_with_cuda(),
-                    "triangular solve op support GPU only now.")
-class TestTriangularSolveOpZeroBatchDim1(TestTriangularSolveOp):
-    def init_case(self):
-        self.inputs = {
-            "input1": np.random.random((3, 3)).astype(np.float32),
-            "input2": np.random.random((3, 1)).astype(np.float32),
-        }
-        self.left_side = True
-        self.upper = True
-        self.transpose_a = False
-        self.unit_diagonal = False
-
-
-@OpTestTool.skip_if(not is_compiled_with_cuda(),
-                    "triangular solve op support GPU only now.")
-class TestTriangularSolveOpZeroBatchDim2(TestTriangularSolveOp):
-    def init_case(self):
-        self.inputs = {
-            "input1": np.random.random((1, 3, 3)).astype(np.float32),
-            "input2": np.random.random((3, 1)).astype(np.float32),
-        }
-        self.left_side = True
-        self.upper = True
-        self.transpose_a = False
-        self.unit_diagonal = False
-
-
-@OpTestTool.skip_if(not is_compiled_with_cuda(),
-                    "triangular solve op support GPU only now.")
-class TestTriangularSolveOpZeroBatchDim3(TestTriangularSolveOp):
-    def init_case(self):
-        self.inputs = {
-            "input1": np.random.random((3, 3)).astype(np.float32),
-            "input2": np.random.random((1, 3, 1)).astype(np.float32),
-        }
-        self.left_side = True
-        self.upper = True
-        self.transpose_a = False
-        self.unit_diagonal = False
+class TestTriangularSolveOpShapeTest(TestCaseHelper):
+    def init_attrs(self):
+        self.class_name = "TestTriangularSolveOpShapeTest"
+        self.cls = TestTriangularSolveOp
+        self.inputs = [
+            {
+                "shape1": [1, 3, 3],
+                "shape2": [1, 3, 1],
+            },
+            {
+                "shape1": [1, 1, 1],
+                "shape2": [1, 1, 1],
+            },
+            {
+                "shape1": [2, 3, 3],
+                "shape2": [2, 3, 10],
+            },
+        ]
+        self.dtypes = [
+            {
+                "dtype": "float32"
+            },
+        ]
+        self.attrs = [
+            {
+                "left_side": True,
+                "upper": True,
+                "transpose_a": False,
+                "unit_diagonal": False
+            },
+        ]
+
+
+class TestTriangularSolveOpDtypeTest(TestCaseHelper):
+    def init_attrs(self):
+        self.class_name = "TestTriangularSolveOpDtypeTest"
+        self.cls = TestTriangularSolveOp
+        self.inputs = [
+            {
+                "shape1": [2, 8, 8],
+                "shape2": [2, 8, 1],
+            },
+        ]
+        self.dtypes = [
+            {
+                "dtype": "float32"
+            },
+            {
+                "dtype": "float64"
+            },
+        ]
+        self.attrs = [
+            {
+                "left_side": True,
+                "upper": True,
+                "transpose_a": False,
+                "unit_diagonal": False
+            },
+        ]
+
+
+class TestTriangularSolveOpBatchDimTest(TestCaseHelper):
+    def init_attrs(self):
+        self.class_name = "TestTriangularSolveOpBatchDimTest"
+        self.cls = TestTriangularSolveOp
+        self.inputs = [
+            {
+                "shape1": [8, 8],
+                "shape2": [8, 4],
+            },
+            {
+                "shape1": [3, 16, 16],
+                "shape2": [16, 4],
+            },
+            {
+                "shape1": [16, 16],
+                "shape2": [5, 16, 4],
+            },
+            {
+                "shape1": [5, 16, 16],
+                "shape2": [5, 16, 4],
+            },
+        ]
+        self.dtypes = [
+            {
+                "dtype": "float32"
+            },
+        ]
+        self.attrs = [
+            {
+                "left_side": True,
+                "upper": True,
+                "transpose_a": False,
+                "unit_diagonal": False
+            },
+        ]
+
+
+class TestTriangularSolveOpBroadcastTest(TestCaseHelper):
+    def init_attrs(self):
+        self.class_name = "TestTriangularSolveOpBroadcastTest"
+        self.cls = TestTriangularSolveOp
+        self.inputs = [
+            {
+                "shape1": [2, 2, 3, 3, 3],
+                "shape2": [1, 3, 4],
+            },
+            {
+                "shape1": [3, 3, 3],
+                "shape2": [2, 2, 3, 3, 4],
+            },
+            {
+                "shape1": [2, 1, 3, 3, 3],
+                "shape2": [2, 2, 3, 3, 4],
+            },
+            {
+                "shape1": [5, 1, 3, 3, 3],
+                "shape2": [1, 2, 1, 3, 4],
+            },
+        ]
+        self.dtypes = [
+            {
+                "dtype": "float32"
+            },
+        ]
+        self.attrs = [
+            {
+                "left_side": True,
+                "upper": True,
+                "transpose_a": False,
+                "unit_diagonal": False
+            },
+        ]
+
+
+class TestTriangularSolveOpAttributeTest(TestCaseHelper):
+    def init_attrs(self):
+        self.class_name = "TestTriangularSolveOpAttributeTest"
+        self.cls = TestTriangularSolveOp
+        self.inputs = [
+            {
+                "shape1": [1, 3, 3],
+                "shape2": [1, 3, 1],
+            },
+            {
+                "shape1": [2, 3, 3],
+                "shape2": [2, 3, 10],
+            },
+        ]
+        self.dtypes = [
+            {
+                "dtype": "float32"
+            },
+        ]
+        self.attrs = [
+            {
+                "left_side": True,
+                "upper": True,
+                "transpose_a": False,
+                "unit_diagonal": False
+            },
+            {
+                "left_side": True,
+                "upper": True,
+                "transpose_a": False,
+                "unit_diagonal": True
+            },
+            {
+                "left_side": True,
+                "upper": True,
+                "transpose_a": True,
+                "unit_diagonal": False
+            },
+        ]
 
 
 @OpTestTool.skip_if(not is_compiled_with_cuda(),
                     "triangular solve op support GPU only now.")
-class TestTriangularSolveOpBroadCast(TestTriangularSolveOp):
-    def init_case(self):
-        self.inputs = {
-            "input1": np.random.random((2, 2, 3, 3, 3)).astype(np.float32),
-            "input2": np.random.random((1, 3, 4)).astype(np.float32),
-        }
-        self.left_side = True
-        self.upper = True
-        self.transpose_a = False
-        self.unit_diagonal = False
-
+class TestTriangularSolveOpRightSide(TestTriangularSolveOp):
+    def setUp(self):
+        self.inputs = {}
+        self.prepare_inputs()
 
-@OpTestTool.skip_if(not is_compiled_with_cuda(),
-                    "triangular solve op support GPU only now.")
-class TestTriangularSolveOpBroadCast1(TestTriangularSolveOp):
-    def init_case(self):
+    def prepare_inputs(self):
         self.inputs = {
-            "input1": np.random.random((3, 3, 3)).astype(np.float32),
-            "input2": np.random.random((2, 2, 3, 3, 4)).astype(np.float32),
+            "input1": self.random([2, 3, 3], "float32"),
+            "input2": self.random([2, 1, 3], "float32")
         }
-        self.left_side = True
+        self.left_side = False
         self.upper = True
         self.transpose_a = False
         self.unit_diagonal = False
 
 
-@OpTestTool.skip_if(not is_compiled_with_cuda(),
-                    "triangular solve op support GPU only now.")
-class TestTriangularSolveOpBroadCast2(TestTriangularSolveOp):
-    def init_case(self):
+class TestTriangularSolveOpRightSide1(TestTriangularSolveOpRightSide):
+    def prepare_inputs(self):
         self.inputs = {
-            "input1": np.random.random((2, 1, 3, 3, 3)).astype(np.float32),
-            "input2": np.random.random((2, 2, 3, 3, 4)).astype(np.float32),
+            "input1": self.random([1, 3, 2, 3, 3], "float32"),
+            "input2": self.random([2, 1, 2, 1, 3], "float32")
         }
-        self.left_side = True
+        self.left_side = False
         self.upper = True
         self.transpose_a = False
         self.unit_diagonal = False
@@ -214,68 +309,35 @@ def init_case(self):
 
 @OpTestTool.skip_if(not is_compiled_with_cuda(),
                     "triangular solve op support GPU only now.")
-class TestTriangularSolveOpBroadCast3(TestTriangularSolveOp):
-    def init_case(self):
-        self.inputs = {
-            "input1": np.random.random((5, 1, 3, 3, 3)).astype(np.float32),
-            "input2": np.random.random((1, 2, 1, 3, 4)).astype(np.float32),
-        }
-        self.left_side = True
-        self.upper = True
-        self.transpose_a = False
-        self.unit_diagonal = False
-
+class TestTriangularSolveOpSingular(TestTriangularSolveOp):
+    def setUp(self):
+        self.inputs = {}
+        self.prepare_inputs()
 
-@OpTestTool.skip_if(not is_compiled_with_cuda(),
-                    "triangular solve op support GPU only now.")
-class TestTriangularSolveOpTranspose(TestTriangularSolveOp):
-    def init_case(self):
+    def prepare_inputs(self):
         self.inputs = {
-            "input1": np.random.random((1, 3, 3)).astype(np.float32),
-            "input2": np.random.random((1, 3, 1)).astype(np.float32),
+            "input1": self.random([1, 3, 3], "float32"),
+            "input2": self.random([1, 3, 1], "float32")
         }
+        # set one dim to zeros to make a singular matrix
+        self.inputs["input1"][0][0] = 0
         self.left_side = True
         self.upper = True
         self.transpose_a = True
         self.unit_diagonal = False
 
-
-@OpTestTool.skip_if(not is_compiled_with_cuda(),
-                    "triangular solve op support GPU only now.")
-class TestTriangularSolveOpRightSide(TestTriangularSolveOp):
-    def init_case(self):
-        self.inputs = {
-            "input1": np.random.random((2, 3, 3)).astype(np.float32),
-            "input2": np.random.random((2, 1, 3)).astype(np.float32),
-        }
-        self.left_side = False
-        self.upper = True
-        self.transpose_a = False
-        self.unit_diagonal = False
-
-
-@OpTestTool.skip_if(not is_compiled_with_cuda(),
-                    "triangular solve op support GPU only now.")
-class TestTriangularSolveOpRightSide1(TestTriangularSolveOp):
-    def init_case(self):
-        self.inputs = {
-            "input1": np.random.random((1, 3, 2, 3, 3)).astype(np.float32),
-            "input2": np.random.random((2, 1, 2, 1, 3)).astype(np.float32),
-        }
-        self.left_side = False
-        self.upper = True
-        self.transpose_a = False
-        self.unit_diagonal = False
+    def test_check_results(self):
+        self.check_outputs_and_grads(equal_nan=True)
 
 
-@OpTestTool.skip_if(not is_compiled_with_cuda(),
-                    "triangular solve op support GPU only now.")
-class TestTriangularSolveOpDoubleFloat(TestTriangularSolveOp):
-    def init_case(self):
+class TestTriangularSolveOpSingular1(TestTriangularSolveOpSingular):
+    def prepare_inputs(self):
         self.inputs = {
-            "input1": np.random.random((1, 3, 3)).astype(np.float64),
-            "input2": np.random.random((1, 3, 1)).astype(np.float64),
+            "input1": self.random([1, 3, 3], "float32"),
+            "input2": self.random([1, 3, 1], "float32")
         }
+        # set one dim to zeros to make a singular matrix
+        self.inputs["input1"][0][2] = 0
         self.left_side = True
         self.upper = True
         self.transpose_a = True
@@ -284,69 +346,47 @@ def init_case(self):
 
 @OpTestTool.skip_if(not is_compiled_with_cuda(),
                     "triangular solve op support GPU only now.")
-class TestTriangularSolveOpBatch(TestTriangularSolveOp):
-    def init_case(self):
-        self.inputs = {
-            "input1": np.random.random((5, 3, 3)).astype(np.float32),
-            "input2": np.random.random((5, 3, 1)).astype(np.float32),
-        }
-        self.left_side = True
-        self.upper = True
-        self.transpose_a = False
-        self.unit_diagonal = False
-
+class TestTriangularSolveOpLarge(TestTriangularSolveOp):
+    def setUp(self):
+        self.inputs = {}
+        self.prepare_inputs()
 
-@OpTestTool.skip_if(not is_compiled_with_cuda(),
-                    "triangular solve op support GPU only now.")
-class TestTriangularSolveOpMultipleRightHandSides(TestTriangularSolveOp):
-    def init_case(self):
+    def prepare_inputs(self):
         self.inputs = {
-            "input1": np.random.random((2, 3, 3)).astype(np.float32),
-            "input2": np.random.random((2, 3, 10)).astype(np.float32),
+            "input1": self.random([1, 1024, 1024], "float64", -0.01, 0.01),
+            "input2": self.random([1, 1024, 512], "float64", -0.01, 0.01)
         }
         self.left_side = True
         self.upper = True
         self.transpose_a = False
         self.unit_diagonal = False
 
-
-@OpTestTool.skip_if(not is_compiled_with_cuda(),
-                    "triangular solve op support GPU only now.")
-class TestTriangularSolveOpSingular(TestTriangularSolveOp):
-    def init_case(self):
-        self.inputs = {
-            "input1": np.random.random((1, 3, 3)).astype(np.float32),
-            "input2": np.random.random((1, 3, 1)).astype(np.float32),
-        }
-        # set one dim to zeros to make a singular matrix
-        self.inputs["input1"][0][0] = 0
-        self.left_side = True
-        self.upper = True
-        self.transpose_a = True
-        self.unit_diagonal = False
-
     def test_check_results(self):
         self.check_outputs_and_grads(equal_nan=True)
 
 
-@OpTestTool.skip_if(not is_compiled_with_cuda(),
-                    "triangular solve op support GPU only now.")
-class TestTriangularSolveOpSingular1(TestTriangularSolveOp):
-    def init_case(self):
+class TestTriangularSolveOpLarge1(TestTriangularSolveOpLarge):
+    def prepare_inputs(self):
         self.inputs = {
-            "input1": np.random.random((1, 3, 3)).astype(np.float32),
-            "input2": np.random.random((1, 3, 1)).astype(np.float32),
+            "input1": self.random([1, 2048, 2048], "float64", -0.01, 0.01),
+            "input2": self.random([1, 2048, 512], "float64", -0.01, 0.01)
         }
-        # set one dim to zeros to make a singular matrix
-        self.inputs["input1"][0][2] = 0
         self.left_side = True
         self.upper = True
-        self.transpose_a = True
+        self.transpose_a = False
         self.unit_diagonal = False
 
-    def test_check_results(self):
-        self.check_outputs_and_grads(equal_nan=True)
-
 
 if __name__ == "__main__":
-    unittest.main()
+    TestTriangularSolveOpShapeTest().run()
+    TestTriangularSolveOpDtypeTest().run()
+    TestTriangularSolveOpBatchDimTest().run()
+    TestTriangularSolveOpBroadcastTest().run()
+    TestTriangularSolveOpAttributeTest().run()
+
+    run_test(TestTriangularSolveOpRightSide)
+    run_test(TestTriangularSolveOpRightSide1)
+    run_test(TestTriangularSolveOpSingular)
+    run_test(TestTriangularSolveOpSingular1)
+    run_test(TestTriangularSolveOpLarge)
+    run_test(TestTriangularSolveOpLarge1)