def test_scope_passdown_when_new_var_created_under_graph_pass():
        """
        If a new_var is created by a graph pass, and the _replace_var happens under the same graph pass,
        the scope information from the old_var is passed to new_var.
        """

        def get_prog():
            @mb.program(input_specs=[mb.TensorSpec(shape=(2, 4))])
            def prog(x):
                with mb.scope(
                    ScopeInfo(source=ScopeSource.TORCHSCRIPT_MODULE_TYPE, data=["module_1"]),
                    ScopeInfo(source=ScopeSource.COREMLTOOLS_GRAPH_PASS, data=["pass_1"]),
                ):
                    # This op is created by pass_1
                    relu = mb.relu(x=x)

                with mb.scope(
                    ScopeInfo(source=ScopeSource.COREMLTOOLS_GRAPH_PASS, data=["pass_2"]),
                ):
                    # This op is newly created by a pass_2
                    sin = mb.sin(x=x)

                with mb.scope(
                    ScopeInfo(source=ScopeSource.TORCHSCRIPT_MODULE_TYPE, data=["module_3"]),
                ):
                    return mb.relu(x=relu)

            return prog

        # Case 1: _replace_var happens outside the graph pass. Nothing happens
        prog = get_prog()
        block = prog.functions["main"]
        ops = list(block.operations)
        var_1, var_2 = ops[0].outputs[0], ops[1].outputs[0]
        assert var_1.scopes == {
            ScopeSource.TORCHSCRIPT_MODULE_TYPE: ["module_1"],
            ScopeSource.COREMLTOOLS_GRAPH_PASS: ["pass_1"],
        }
        assert var_2.scopes == {
            ScopeSource.COREMLTOOLS_GRAPH_PASS: ["pass_2"],
        }

        block._replace_var(var_1, var_2)
        assert var_1.scopes == {
            ScopeSource.TORCHSCRIPT_MODULE_TYPE: ["module_1"],
            ScopeSource.COREMLTOOLS_GRAPH_PASS: ["pass_1"],
        }
        assert var_2.scopes == {
            ScopeSource.COREMLTOOLS_GRAPH_PASS: ["pass_2"],
        }

        # Case 2: new_var created under a pass_2, and _replace_var happens under pass_2. Scope info is passed from the old_var
        # to the new_var
        @mb.program(input_specs=[mb.TensorSpec(shape=(2, 4))])
        def prog(x):
            with mb.scope(
                ScopeInfo(source=ScopeSource.TORCHSCRIPT_MODULE_TYPE, data=["module_1"]),
                ScopeInfo(source=ScopeSource.COREMLTOOLS_GRAPH_PASS, data=["pass_1"]),
            ):
                # This op is created by pass_1
                relu = mb.relu(x=x)

            with mb.scope(
                ScopeInfo(source=ScopeSource.TORCHSCRIPT_MODULE_TYPE, data=["module_3"]),
            ):
                return mb.relu(x=relu)

        with prog.functions["main"] as block:
            op_1, op_2 = list(block.operations)
            with mb.scope(
                ScopeInfo(source=ScopeSource.COREMLTOOLS_GRAPH_PASS, data=["pass_2"]),
            ):
                # This op is newly created by a pass_2
                sin = mb.sin(x=block.inputs["x"], before_op=op_2)
                block._replace_var(op_1.outputs[0], sin)

        block = prog.functions["main"]
        ops = list(block.operations)
        var_1, var_2 = ops[0].outputs[0], ops[1].outputs[0]

        assert var_1.scopes == {
            ScopeSource.TORCHSCRIPT_MODULE_TYPE: ["module_1"],
            ScopeSource.COREMLTOOLS_GRAPH_PASS: ["pass_1"],
        }
        assert var_2.scopes == {
            ScopeSource.TORCHSCRIPT_MODULE_TYPE: ["module_1"],
            ScopeSource.COREMLTOOLS_GRAPH_PASS: ["pass_1", "pass_2"],
        }

        # Case 3: new_var created under a pass_2, but _replace_var happens under pass_3.
        # Nothing happens.
        prog = get_prog()
        block = prog.functions["main"]
        ops = list(block.operations)
        var_1, var_2 = ops[0].outputs[0], ops[1].outputs[0]
        with mb.scope(ScopeInfo(source=ScopeSource.COREMLTOOLS_GRAPH_PASS, data=["pass_3"])):
            block._replace_var(var_1, var_2)
        assert var_1.scopes == {
            ScopeSource.TORCHSCRIPT_MODULE_TYPE: ["module_1"],
            ScopeSource.COREMLTOOLS_GRAPH_PASS: ["pass_1"],
        }
        assert var_2.scopes == {
            ScopeSource.COREMLTOOLS_GRAPH_PASS: ["pass_2"],
        }

        # Case 4: new_var created under pass_2, and be passed down some scope info,
        # so even though _replace_var happens under pass_2 again, nothing happens.
        @mb.program(input_specs=[mb.TensorSpec(shape=(2, 4))])
        def prog(x):
            with mb.scope(
                ScopeInfo(source=ScopeSource.TORCHSCRIPT_MODULE_TYPE, data=["module_1"]),
                ScopeInfo(source=ScopeSource.COREMLTOOLS_GRAPH_PASS, data=["pass_1"]),
            ):
                # This op is created by pass_1
                relu = mb.relu(x=x)

            with mb.scope(
                ScopeInfo(source=ScopeSource.COREMLTOOLS_GRAPH_PASS, data=["pass_2"]),
                ScopeInfo(source=ScopeSource.TORCHSCRIPT_MODULE_TYPE, data=["module_2"]),
            ):
                # This op is newly created by a pass_2, and other scope info already passed down
                sin = mb.sin(x=x)

            with mb.scope(
                ScopeInfo(source=ScopeSource.TORCHSCRIPT_MODULE_TYPE, data=["module_3"]),
            ):
                return mb.relu(x=relu)

        block = prog.functions["main"]
        ops = list(block.operations)
        var_1, var_2 = ops[0].outputs[0], ops[1].outputs[0]
        assert var_1.scopes == {
            ScopeSource.TORCHSCRIPT_MODULE_TYPE: ["module_1"],
            ScopeSource.COREMLTOOLS_GRAPH_PASS: ["pass_1"],
        }
        assert var_2.scopes == {
            ScopeSource.TORCHSCRIPT_MODULE_TYPE: ["module_2"],
            ScopeSource.COREMLTOOLS_GRAPH_PASS: ["pass_2"],
        }

        with mb.scope(ScopeInfo(source=ScopeSource.COREMLTOOLS_GRAPH_PASS, data=["pass_2"])):
            block._replace_var(var_1, var_2)
        assert var_1.scopes == {
            ScopeSource.TORCHSCRIPT_MODULE_TYPE: ["module_1"],
            ScopeSource.COREMLTOOLS_GRAPH_PASS: ["pass_1"],
        }
        assert var_2.scopes == {
            ScopeSource.TORCHSCRIPT_MODULE_TYPE: ["module_2"],
            ScopeSource.COREMLTOOLS_GRAPH_PASS: ["pass_2"],
        }

        # Case 5: new_var created under pass_2, but the graph pass already finished,
        # so even though _replace_var happens under pass_2 again, nothing happens.
        @mb.program(input_specs=[mb.TensorSpec(shape=(2, 4))])
        def prog(x):
            with mb.scope(
                ScopeInfo(source=ScopeSource.COREMLTOOLS_GRAPH_PASS, data=["pass_1"]),
            ):
                # This op is created by pass_1
                relu = mb.relu(x=x)

            with mb.scope(
                ScopeInfo(source=ScopeSource.COREMLTOOLS_GRAPH_PASS, data=["pass_2"]),
            ):
                # This op is newly created by a pass_2, and other scope info already passed down
                sin = mb.sin(x=x)

            with mb.scope(
                ScopeInfo(source=ScopeSource.TORCHSCRIPT_MODULE_TYPE, data=["module_3"]),
            ):
                return mb.relu(x=relu)

        block = prog.functions["main"]
        ops = list(block.operations)
        var_1, var_2 = ops[0].outputs[0], ops[1].outputs[0]
        assert var_1.scopes == {
            ScopeSource.COREMLTOOLS_GRAPH_PASS: ["pass_1"],
        }
        assert var_2.scopes == {
            ScopeSource.COREMLTOOLS_GRAPH_PASS: ["pass_2"],
        }

        with mb.scope(ScopeInfo(source=ScopeSource.COREMLTOOLS_GRAPH_PASS, data=["pass_2"])):
            block._replace_var(var_1, var_2)
        assert var_1.scopes == {
            ScopeSource.COREMLTOOLS_GRAPH_PASS: ["pass_1"],
        }
        assert var_2.scopes == {
            ScopeSource.COREMLTOOLS_GRAPH_PASS: ["pass_2"],
        }

        # Case 6: new_var created under nested graph passes scope. And graph pass happens under pass_3.
        @mb.program(input_specs=[mb.TensorSpec(shape=(2, 4))])
        def prog(x):
            with mb.scope(
                ScopeInfo(source=ScopeSource.TORCHSCRIPT_MODULE_TYPE, data=["module_1"]),
                ScopeInfo(source=ScopeSource.COREMLTOOLS_GRAPH_PASS, data=["pass_1"]),
            ):
                # This op is created by pass_1
                relu = mb.relu(x=x)

            with mb.scope(
                ScopeInfo(source=ScopeSource.TORCHSCRIPT_MODULE_TYPE, data=["module_3"]),
            ):
                return mb.relu(x=relu)

        block = prog.functions["main"]
        ops = list(block.operations)

        with block:
            with mb.scope(
                ScopeInfo(source=ScopeSource.COREMLTOOLS_GRAPH_PASS, data=["pass_2"]),
            ):
                with mb.scope(
                    ScopeInfo(source=ScopeSource.COREMLTOOLS_GRAPH_PASS, data=["pass_3"]),
                ):
                    sin = mb.sin(x=block.inputs["x"], before_op=ops[1])
                    block._replace_var(ops[0].outputs[0], sin)

        ops = list(block.operations)
        var_1, var_2 = ops[0].outputs[0], ops[1].outputs[0]
        assert var_1.scopes == {
            ScopeSource.TORCHSCRIPT_MODULE_TYPE: ["module_1"],
            ScopeSource.COREMLTOOLS_GRAPH_PASS: ["pass_1"],
        }
        assert var_2.scopes == {
            ScopeSource.TORCHSCRIPT_MODULE_TYPE: ["module_1"],
            ScopeSource.COREMLTOOLS_GRAPH_PASS: ["pass_1", "pass_2", "pass_3"],
        }

        # Case 7: new_var created under nested graph passes scope. And graph pass happens under pass_2. Nothing will happen in this case, since new_var is created under pass_3.
        @mb.program(input_specs=[mb.TensorSpec(shape=(2, 4))])
        def prog(x):
            with mb.scope(
                ScopeInfo(source=ScopeSource.TORCHSCRIPT_MODULE_TYPE, data=["module_1"]),
                ScopeInfo(source=ScopeSource.COREMLTOOLS_GRAPH_PASS, data=["pass_1"]),
            ):
                # This op is created by pass_1
                relu = mb.relu(x=x)

            with mb.scope(
                ScopeInfo(source=ScopeSource.TORCHSCRIPT_MODULE_TYPE, data=["module_3"]),
            ):
                return mb.relu(x=relu)

        block = prog.functions["main"]
        ops = list(block.operations)

        with block:
            with mb.scope(
                ScopeInfo(source=ScopeSource.COREMLTOOLS_GRAPH_PASS, data=["pass_2"]),
            ):
                with mb.scope(
                    ScopeInfo(source=ScopeSource.COREMLTOOLS_GRAPH_PASS, data=["pass_3"]),
                ):
                    sin = mb.sin(x=block.inputs["x"], before_op=ops[1])
                block._replace_var(ops[0].outputs[0], sin)

        ops = list(block.operations)
        var_1, var_2 = ops[0].outputs[0], ops[1].outputs[0]
        assert var_1.scopes == {
            ScopeSource.TORCHSCRIPT_MODULE_TYPE: ["module_1"],
            ScopeSource.COREMLTOOLS_GRAPH_PASS: ["pass_1"],
        }
        assert var_2.scopes == {
            ScopeSource.COREMLTOOLS_GRAPH_PASS: ["pass_2", "pass_3"],
        }

        If a new_var is created by a graph pass, and the _replace_var happens under the same graph pass,
        the scope information from the old_var is passed to new_var.
        
coremltools/converters/mil/mil/tests/test_programs.py