def test_time_complexity(self):
        """
        This test makes sure the cast_optimization's time complexity is O(N) for most of the cases.

        In this test case, the program consists of 1000 relu ops followed by 100 cast ops.

            input -> relu -> relu -> ... -> relu -> cast -> cast -> ... -> cast

        The algorithm goes through the first pass to eliminate all cast ops:

            input -> relu -> ... -> relu

        Note that, the total number of visited op is 1000 (relu) + 100 (cast) + 100 (const for the dtype) = 1200.

        Because the fusion happens, the algorithm goes through the program again.
        This time, the number of visited op is 1000 (relu) + 100 (const) = 1100.

        Overally, the number of visited op is 2300.
        """

        @mb.program(input_specs=[mb.TensorSpec(shape=(1,), dtype=types.fp32)])
        def prog(x):
            for _ in range(1000):
                x = mb.relu(x=x)
            for _ in range(100):
                x = mb.cast(x=x, dtype="fp32")
            return x

        graph_pass = cast_optimization()
        graph_pass.apply(prog)
        assert (
            graph_pass._num_of_visited_ops == 2_300
        )  # Please refer to the doc string for how 2300 comes from.

        This test makes sure the cast_optimization's time complexity is O(N) for most of the cases.

        In this test case, the program consists of 1000 relu ops followed by 100 cast ops.

            input -> relu -> relu -> ... -> relu -> cast -> cast -> ... -> cast

        The algorithm goes through the first pass to eliminate all cast ops:

            input -> relu -> ... -> relu

        Note that, the total number of visited op is 1000 (relu) + 100 (cast) + 100 (const for the dtype) = 1200.

        Because the fusion happens, the algorithm goes through the program again.
        This time, the number of visited op is 1000 (relu) + 100 (const) = 1100.

        Overally, the number of visited op is 2300.
        
coremltools/converters/mil/mil/passes/tests/test_passes.py