def test_simple():
        """
        Given:

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

        Return:

            mul = mul(state, x)
            update = coreml_update_state(state, mul)
            add = add(update, y)
        """
        SHAPE = (2, 3, 5)

        @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=x, y=read)
            update = mb.coreml_update_state(state=state, value=mul)
            add = mb.add(x=x, y=mul)
            return add

        mul_op = prog.find_ops(op_type="mul")[0]
        add_op = prog.find_ops(op_type="add")[0]
        assert add_op.y is mul_op.outputs[0]

        apply_pass_and_basic_check(prog, "common::prefer_state_in_downstream")
        coreml_update_state_op = prog.find_ops(op_type="coreml_update_state")[0]
        add_op = prog.find_ops(op_type="add")[0]
        assert add_op.y is coreml_update_state_op.outputs[0]

        Given:

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

        Return:

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