def test_linear_consecutive_cast_ops_cancellation(self):
        """Test the cast optimization pass with more complicated patterns."""

        """
        Input graph:
        input(fp16) -----> cast(dtype="fp32") -----> cast(dtype="fp16") ----> square ---> out

        Output graph:
        input -----> square -----> out
        """
        @mb.program(input_specs=[mb.TensorSpec(shape=(10, 20), dtype=types.fp16)])
        def prog(x):
            x = mb.cast(x=x, dtype="fp32")
            x = mb.cast(x=x, dtype="fp16")
            x = mb.square(x=x)
            return x

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

        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) == ["square"]

        assert_model_is_valid(
            prog,
            {"x": (10, 20)},
            expected_output_shapes={block.outputs[0].name: (10, 20)},
        )
Test the cast optimization pass with more complicated patterns.
coremltools/converters/mil/mil/passes/tests/test_passes.py