def test_coreml_update_state_lowering():
        """
        If the output of coreml_update_state is not a block output and
        it is not fed into any other ops, the op should be translated into
        a single write_state.
        """

        @mb.program(
            input_specs=[
                mb.StateTensorSpec((1,), dtype=types.fp16),
                mb.TensorSpec((1,), dtype=types.fp16),
                mb.TensorSpec((1,), dtype=types.fp16),
            ],
            opset_version=ct.target.iOS18,
        )
        def prog(state, x, y):
            mb.coreml_update_state(state=state, value=x)
            mb.coreml_update_state(state=state, value=y)
            return x, mb.coreml_update_state(state=state, value=y)

        mlmodel = ct.convert(
            prog,
            convert_to="mlprogram",
            minimum_deployment_target=ct.target.iOS18,
            skip_model_load=True,
        )

        mil = mlmodel.get_spec().mlProgram

        for function in mil.functions.values():
            for block in function.block_specializations.values():
                ops = list(block.operations)
                expected_ops = [
                    "write_state",
                    "write_state",
                    "write_state",
                    "read_state",
                ]
                assert [val.type for val in ops] == expected_ops

        If the output of coreml_update_state is not a block output and
        it is not fed into any other ops, the op should be translated into
        a single write_state.
        
coremltools/converters/mil/backend/mil/test_load.py