def test_loop_invariant_elimination2(self):
        """
        Invariant pattern: Block outputs var from outside of the block
        """

        @mb.program(
            input_specs=[
                mb.TensorSpec(shape=(1, 2)),
                mb.TensorSpec(shape=(1, 2)),
            ]
        )
        def prog(a, b):
            def body(a, bx):
                return mb.add(x=a, y=b), b

            def cond(a, bx):
                a_mean = mb.reduce_mean(x=a, axes=[0, 1])
                b_mean = mb.reduce_mean(x=bx, axes=[0, 1])
                return mb.less(x=a_mean, y=b_mean)

            # b is loop invariant
            return mb.while_loop(_cond=cond, _body=body, loop_vars=(a, b))

        while_op = prog.find_ops(op_type="while_loop", exactly_one=True)[0]
        assert len(while_op.blocks[0].inputs) == 2
        assert len(while_op.outputs) == 2
        assert len(while_op.loop_vars) == 2
        assert while_op.blocks[0].inputs[0].name == "a_x0"
        assert while_op.blocks[0].inputs[1].name == "b_x0"

        prev_prog = copy.deepcopy(prog)
        PASS_REGISTRY["common::loop_invariant_elimination"](prog)
        assert_same_output_names(prev_prog, prog)

        while_op = prog.find_ops(op_type="while_loop", exactly_one=True)[0]
        assert len(while_op.blocks[0].inputs) == 1
        assert len(while_op.outputs) == 1
        assert len(while_op.loop_vars) == 1
        assert while_op.blocks[0].inputs[0].name == "a_x0"

        if _VALIDATE_MODEL:
            assert_model_is_valid(prog, {"a": (1, 2), "b": (1, 2)})

        Invariant pattern: Block outputs var from outside of the block
        
coremltools/converters/mil/mil/passes/tests/test_cleanup_passes.py