def test_fp16_scope_preservation(self):
        """
        This test explains step-by-step how the scope information preservation works in the fp32 -> fp16 pass.

        Input graph:

            x(fp32)
            -> relu(torch_scope="module_1")
            -> sin(torch_scope="module_2")
            -> output(fp32)

        (1) "common::add_fp16_cast"

            First, in the add_fp16_cast graph pass, multiple cast ops are injected in the graph:

                x(fp32)
                -> cast(dtype="fp16", torch_scope="module_1", pass_scope="add_fp16_cast")
                -> relu(torch_scope="module_1", pass_scope="add_fp16_cast")
                -> cast(dtype="fp32", torch_scope="module_1", pass_scope="add_fp16_cast")
                -> cast(dtype="fp16", torch_scope="module_2", pass_scope="add_fp16_cast")
                -> sin(torch_scope="module_2", pass_scope="add_fp16_cast")
                -> cast(dtype="fp32, torch_scope="module_2", pass_scope="add_fp16_cast")
                -> output

            There are 4 cast ops in the graph who has pass_scope = "add_fp16_cast", which indicates they are added by the "add_fp16_cast" pass.

            Note that, the first cast -> relu -> cast pattern has the same torch scope information as
            the original relu(torch_scope="module_1"). This is due to the fact that when we replace
            the use of the original relu output with the output of the second cast op, the scope information is back propagated.

            The same reason applied for why the cast -> sin -> cast patterns has the torch scope as
            the original sin op.

        (2) "common::cast_optimization" + "dead-code_elimination"

            After the cleanup, the graph becomes:

                x(fp32)
                -> cast(
                        dtype="fp16",
                        torch_scope="module_1",
                        pass_scope="add_fp16_cast"
                    )
                -> relu(
                        torch_scope="module_1",
                        pass_scope="add_fp16_cast]
                    )
                -> sin(
                        torch_scope="module_2",
                        pass_scope="add_fp16_cast"
                    )
                -> cast(
                        dtype="fp32,
                        torch_scope="module_2",
                        pass_scope="add_fp16_cast"
                    )
                -> output

            We can see that, the fp16 version of relu / sin preserves the original torch scope information.
        """
        shape = (3, 5, 6, 7)

        @mb.program(input_specs=[mb.TensorSpec(shape=shape)])
        def prog(x):
            with mb.scope(
                ScopeInfo(source=ScopeSource.TORCHSCRIPT_MODULE_TYPE, data="module_1"),
            ):
                x = mb.relu(x=x)
            with mb.scope(
                ScopeInfo(source=ScopeSource.TORCHSCRIPT_MODULE_TYPE, data="module_2"),
            ):
                return mb.sin(x=x)

        prog._add_essential_scope_source(ScopeSource.TORCHSCRIPT_MODULE_TYPE)

        # fp16 cast pass
        apply_pass_and_basic_check(prog, "common::add_fp16_cast")
        assert get_op_types_in_program(prog) == ["cast", "relu", "cast", "cast", "sin", "cast"]

        cast_ops = prog.find_ops(op_type="cast")
        assert len(cast_ops) == 4
        assert len(cast_ops[0].scopes) == 2
        assert cast_ops[0].scopes[ScopeSource.COREMLTOOLS_GRAPH_PASS] == ["add_fp16_cast"]
        assert cast_ops[0].scopes[ScopeSource.TORCHSCRIPT_MODULE_TYPE] == ["module_1"]
        assert len(cast_ops[1].scopes) == 2
        assert cast_ops[1].scopes[ScopeSource.COREMLTOOLS_GRAPH_PASS] == ["add_fp16_cast"]
        assert cast_ops[1].scopes[ScopeSource.TORCHSCRIPT_MODULE_TYPE] == ["module_1"]
        assert len(cast_ops[2].scopes) == 2
        assert cast_ops[2].scopes[ScopeSource.COREMLTOOLS_GRAPH_PASS] == [
            "add_fp16_cast",
        ]
        assert cast_ops[2].scopes[ScopeSource.TORCHSCRIPT_MODULE_TYPE] == ["module_2"]
        assert len(cast_ops[3].scopes) == 2
        assert cast_ops[3].scopes[ScopeSource.COREMLTOOLS_GRAPH_PASS] == ["add_fp16_cast"]
        assert cast_ops[3].scopes[ScopeSource.TORCHSCRIPT_MODULE_TYPE] == ["module_2"]

        relu_op = prog.find_ops(op_type="relu")[0]
        assert len(relu_op.scopes) == 2
        assert relu_op.scopes[ScopeSource.COREMLTOOLS_GRAPH_PASS] == ["add_fp16_cast"]
        assert relu_op.scopes[ScopeSource.TORCHSCRIPT_MODULE_TYPE] == ["module_1"]

        sin_op = prog.find_ops(op_type="sin")[0]
        assert len(sin_op.scopes) == 2
        assert sin_op.scopes[ScopeSource.COREMLTOOLS_GRAPH_PASS] == ["add_fp16_cast"]
        assert sin_op.scopes[ScopeSource.TORCHSCRIPT_MODULE_TYPE] == ["module_2"]

        # clean up with cast optimization and dead code elimination
        apply_pass_and_basic_check(prog, "common::cast_optimization")
        apply_pass_and_basic_check(prog, "common::dead_code_elimination")

        assert get_op_types_in_program(prog) == ["cast", "relu", "sin", "cast"]
        cast_ops = prog.find_ops(op_type="cast")
        assert len(cast_ops) == 2
        assert len(cast_ops[0].scopes) == 2
        assert cast_ops[0].scopes[ScopeSource.COREMLTOOLS_GRAPH_PASS] == ["add_fp16_cast"]
        assert cast_ops[0].scopes[ScopeSource.TORCHSCRIPT_MODULE_TYPE] == ["module_1"]
        assert len(cast_ops[1].scopes) == 2
        assert cast_ops[1].scopes[ScopeSource.COREMLTOOLS_GRAPH_PASS] == ["add_fp16_cast"]
        assert cast_ops[1].scopes[ScopeSource.TORCHSCRIPT_MODULE_TYPE] == ["module_2"]

        relu_op = prog.find_ops(op_type="relu")[0]
        assert len(relu_op.scopes) == 2
        assert relu_op.scopes[ScopeSource.COREMLTOOLS_GRAPH_PASS] == [
            "add_fp16_cast",
        ]
        assert relu_op.scopes[ScopeSource.TORCHSCRIPT_MODULE_TYPE] == ["module_1"]

        sin_op = prog.find_ops(op_type="sin")[0]
        assert len(sin_op.scopes) == 2
        assert sin_op.scopes[ScopeSource.COREMLTOOLS_GRAPH_PASS] == ["add_fp16_cast"]
        assert sin_op.scopes[ScopeSource.TORCHSCRIPT_MODULE_TYPE] == ["module_2"]

        This test explains step-by-step how the scope information preservation works in the fp32 -> fp16 pass.

        Input graph:

            x(fp32)
            -> relu(torch_scope="module_1")
            -> sin(torch_scope="module_2")
            -> output(fp32)

        (1) "common::add_fp16_cast"

            First, in the add_fp16_cast graph pass, multiple cast ops are injected in the graph:

                x(fp32)
                -> cast(dtype="fp16", torch_scope="module_1", pass_scope="add_fp16_cast")
                -> relu(torch_scope="module_1", pass_scope="add_fp16_cast")
                -> cast(dtype="fp32", torch_scope="module_1", pass_scope="add_fp16_cast")
                -> cast(dtype="fp16", torch_scope="module_2", pass_scope="add_fp16_cast")
                -> sin(torch_scope="module_2", pass_scope="add_fp16_cast")
                -> cast(dtype="fp32, torch_scope="module_2", pass_scope="add_fp16_cast")
                -> output

            There are 4 cast ops in the graph who has pass_scope = "add_fp16_cast", which indicates they are added by the "add_fp16_cast" pass.

            Note that, the first cast -> relu -> cast pattern has the same torch scope information as
            the original relu(torch_scope="module_1"). This is due to the fact that when we replace
            the use of the original relu output with the output of the second cast op, the scope information is back propagated.

            The same reason applied for why the cast -> sin -> cast patterns has the torch scope as
            the original sin op.

        (2) "common::cast_optimization" + "dead-code_elimination"

            After the cleanup, the graph becomes:

                x(fp32)
                -> cast(
                        dtype="fp16",
                        torch_scope="module_1",
                        pass_scope="add_fp16_cast"
                    )
                -> relu(
                        torch_scope="module_1",
                        pass_scope="add_fp16_cast]
                    )
                -> sin(
                        torch_scope="module_2",
                        pass_scope="add_fp16_cast"
                    )
                -> cast(
                        dtype="fp32,
                        torch_scope="module_2",
                        pass_scope="add_fp16_cast"
                    )
                -> output

            We can see that, the fp16 version of relu / sin preserves the original torch scope information.
        
coremltools/converters/mil/mil/passes/tests/test_passes.py