def test_simple():
        """
        Given:

            mul = mul(state, x)
            add = add(mul, 1.0)
            update = coreml_update_state(state, mul)

        Return:

            mul = mul(state, x)
            update = coreml_update_state(state, mul)
            add = add(mul, 1.0)
        """
        SHAPE = (2,)

        @mb.program(
            input_specs=[
                mb.TensorSpec(shape=SHAPE, dtype=types.fp16),
                mb.StateTensorSpec(shape=SHAPE, dtype=types.fp16),
            ],
            opset_version=AvailableTarget.iOS18,
        )
        def prog(x, state):
            read = mb.read_state(input=state)
            mul = mb.mul(x=read, y=x)
            add = mb.add(x=mul, y=np.float16(1.0))
            update = mb.coreml_update_state(state=state, value=mul)
            return add

        assert get_op_types_in_program(prog) == ["read_state", "mul", "add", "coreml_update_state"]

        apply_pass_and_basic_check(prog, "common::canonicalize_inplace_pattern")
        assert get_op_types_in_program(prog) == ["read_state", "mul", "coreml_update_state", "add"]

        Given:

            mul = mul(state, x)
            add = add(mul, 1.0)
            update = coreml_update_state(state, mul)

        Return:

            mul = mul(state, x)
            update = coreml_update_state(state, mul)
            add = add(mul, 1.0)
        
coremltools/converters/mil/mil/passes/tests/test_state_passes.py