def test_consecutive_fusable_casts_on_all_branches(self):
        """
        Input graph:
                                         |---->cast(dtype="int32")---->square--->out_1
                                         |
        input(fp16)---->cast(dtype="fp32")---->cast(dtype="int32")---->abs--->out_2
                                         |
                                         |---->cast(dtype="int32")---->identity--->out_3

        Output graph:

                                          |-->square-->out_1
                                          |
        input(fp16)---->cast(dtype="int32")-->abs-->out_2
                                          |
                                          |-->identity->out_3

        Note that, this result needs the assistant of another pass remove_redundant_ops
        """
        @mb.program(input_specs=[mb.TensorSpec(shape=(10, 20), dtype=types.fp16)])
        def prog(x):
            x = mb.cast(x=x, dtype="fp32")
            x1 = mb.cast(x=x, dtype="int32")
            x2 = mb.cast(x=x, dtype="int32")
            x3 = mb.cast(x=x, dtype="int32")
            x4 = mb.square(x=x1)
            x5 = mb.abs(x=x2)
            x6 = mb.identity(x=x3)
            return x4, x5, x6

        assert get_op_types_in_program(prog) == [
            "cast",
            "cast",
            "cast",
            "cast",
            "square",
            "abs",
            "identity",
        ]

        apply_pass_and_basic_check(prog, "common::cast_optimization")
        _, _, block = apply_pass_and_basic_check(prog, "common::dead_code_elimination")
        assert get_op_types_in_program(prog) == [
            "cast",
            "cast",
            "cast",
            "square",
            "abs",
            "identity",
        ]
        cast_ops = block.find_ops(op_type="cast")
        assert all(v.dtype.val == "int32" for v in cast_ops)

        apply_pass_and_basic_check(prog, "common::remove_redundant_ops")
        _, _, block = apply_pass_and_basic_check(prog, "common::dead_code_elimination")
        assert get_op_types_in_program(prog) == [
            "cast",
            "square",
            "abs",
            "identity",
        ]
        assert block.find_ops(op_type="cast")[0].dtype.val == "int32"

        assert_model_is_valid(
            prog,
            {"x": (10, 20)},
            expected_output_shapes={
                block.outputs[0].name: (10, 20),
                block.outputs[1].name: (10, 20),
                block.outputs[2].name: (10, 20),
            },
        )

        Input graph:
                                         |---->cast(dtype="int32")---->square--->out_1
                                         |
        input(fp16)---->cast(dtype="fp32")---->cast(dtype="int32")---->abs--->out_2
                                         |
                                         |---->cast(dtype="int32")---->identity--->out_3

        Output graph:

                                          |-->square-->out_1
                                          |
        input(fp16)---->cast(dtype="int32")-->abs-->out_2
                                          |
                                          |-->identity->out_3

        Note that, this result needs the assistant of another pass remove_redundant_ops
        
coremltools/converters/mil/mil/passes/tests/test_passes.py