def test_simple_2():
        """
        Input graph:

            input(fp32) -> identity -> cast(dtype="int32") -> output_1
                               |
                               .-> output_2

        Output graph:

            input(fp16) -> cast(dtype="fp32") -> identity -> cast(dtype="int32")  -> output_1
                                                      |
                                                      .-> output_2,

            with function.output_types = [ct.TesorType(dtype=types.int32), ct.TesorType(dtype=types.fp16)]

        """

        @mb.program(input_specs=[mb.TensorSpec(shape=(2, 3))])
        def prog(x):
            x = mb.identity(x=x)
            return mb.cast(x=x, dtype="int32"), x

        graph_pass = add_fp16_cast()
        block = prog.functions["main"]
        graph_pass.transform_function_signatures(block)

        assert block.inputs["x"].dtype == types.fp16

        assert get_op_types_in_program(prog) == ["cast", "identity", "cast"]
        cast_ops = block.find_ops(op_type="cast")
        assert cast_ops[0].dtype.val == "fp32"
        assert cast_ops[1].dtype.val == "int32"

        assert len(block.outputs) == 2
        assert block.outputs[0].dtype == types.int32
        assert block.outputs[1].dtype == types.fp32

        assert len(block.output_types) == 2
        assert block.output_types[0].dtype == types.int32
        assert block.output_types[1].dtype == types.fp16

        Input graph:

            input(fp32) -> identity -> cast(dtype="int32") -> output_1
                               |
                               .-> output_2

        Output graph:

            input(fp16) -> cast(dtype="fp32") -> identity -> cast(dtype="int32")  -> output_1
                                                      |
                                                      .-> output_2,

            with function.output_types = [ct.TesorType(dtype=types.int32), ct.TesorType(dtype=types.fp16)]

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