def test_different_consecutive_casts_config_on_different_branches(self):
        """
        Input graph:

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


        Output graph:

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

        """
        @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="fp16")
            x2 = mb.cast(x=x, dtype="int32")
            x3 = mb.square(x=x1)
            x4 = mb.exp2(x=x2)
            x5 = mb.abs(x=x)
            return x3, x4, x5

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

        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", "square", "exp2", "abs"]

        # Asserting first cast configuration
        cast_1 = block.find_ops(op_type="cast")[0]
        assert cast_1.dtype.val == "fp32"
        assert len(cast_1.outputs) == 1
        assert len(cast_1.outputs[0].child_ops) == 1
        assert cast_1.outputs[0].child_ops[0].op_type == "abs"

        # Asserting second cast configuration
        cast_2 = block.find_ops(op_type="cast")[1]
        assert cast_2.dtype.val == "int32"
        assert len(cast_2.outputs) == 1
        assert len(cast_2.outputs[0].child_ops) == 1
        assert cast_2.outputs[0].child_ops[0].op_type == "exp2"

        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="fp16")---->square--->out_1
                                        |
        input(fp16)---->cast(dtype="fp32")---->cast(dtype="int32")---->exp2--->out_2
                                        |
                                        |---->abs--->out_3


        Output graph:

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

        
coremltools/converters/mil/mil/passes/tests/test_passes.py