def test_single_state_single_input(
        self, float32_buffer_model, float32_non_persistent_buffer_model, compute_unit
    ):
        """
        Tests for different combination of input dtypes.
        """

        def test_valid_prog(prog, expected_ops=None):
            block = prog.functions["main"]
            assert types.is_tensor(block.inputs["x"].sym_type)
            assert types.is_state(block.inputs["state_workaround"].sym_type)
            assert len(block.outputs) == 1
            assert types.is_tensor(block.outputs[0].sym_type)
            if expected_ops is None:
                expected_ops = [
                    "read_state",
                    "add",
                    "coreml_update_state",
                    "mul",
                ]
            assert get_op_types_in_program(prog) == expected_ops

        """
        fp32 state / input (default with compute_precision=fp32),
        with both persistent and non-persistent buffer.
        fp32 state is not supported through runtime.

        (%x: Tensor(fp32), %state: State(fp32)) -> {
            %read_state(fp32) = read_state(%state)
            %add(fp32) = add(%read_state, %x)
            %update(fp32) = coreml_update_state(%state, %add)
            %mul(fp32) = mul(%update, 5)
        } -> (%mul)
        """
        for model in [float32_buffer_model, float32_non_persistent_buffer_model]:
            prog = ct.convert(
                model,
                inputs=[
                    ct.TensorType(shape=(3,)),
                ],
                states=[
                    ct.StateType(
                        wrapped_type=ct.TensorType(
                            shape=(3,),
                        ),
                        name="state",
                    ),
                ],
                minimum_deployment_target=ct.target.iOS18,
                compute_precision=ct.precision.FLOAT32,
                convert_to="milinternal",
            )
            test_valid_prog(prog)
            block = prog.functions["main"]
            assert block.inputs["x"].sym_type.get_primitive() == types.fp32
            assert (
                block.inputs["state_workaround"].sym_type.wrapped_type().get_primitive()
                == types.fp32
            )
            assert block.outputs[0].dtype == types.fp32

        """
        fp16 state / input (user specify)

        (%x: Tensor(fp16), %state: State(fp16)) -> {
            %read_state(fp16) = read_state(%state)
            %add(fp16) = add(%read_state, %x)
            %update(fp16) = coreml_update_state(%state, %add)
            %mul(fp16) = mul(%update, 5)
        } -> (%mul)
        """
        mlmodel = ct.convert(
            float32_buffer_model,
            inputs=[
                ct.TensorType(shape=(3,), dtype=np.float16),
            ],
            states=[
                ct.StateType(
                    wrapped_type=ct.TensorType(
                        shape=(3,),
                        dtype=np.float16,
                    ),
                    name="state",
                ),
            ],
            minimum_deployment_target=ct.target.iOS18,
            convert_to="mlprogram",
            compute_units=compute_unit,
        )
        # check the pymil program
        prog = mlmodel._mil_program
        test_valid_prog(prog)
        block = prog.functions["main"]
        assert block.inputs["x"].sym_type.get_primitive() == types.fp16
        assert (
            block.inputs["state_workaround"].sym_type.wrapped_type().get_primitive() == types.fp16
        )
        assert block.outputs[0].dtype == types.fp16

        # check the mil proto
        mil = mlmodel.get_spec().mlProgram
        for function in mil.functions.values():
            for block in function.block_specializations.values():
                ops = list(block.operations)
                expected_ops = [
                    "read_state",
                    "add",
                    "write_state",
                    "read_state",
                    "const",
                    "mul",
                ]
                assert [val.type for val in ops] == expected_ops
                assert len(ops[2].outputs) == 0

        verify_prediction(mlmodel)

        """
        fp16 state / input (default with compute_precision=fp16)

        (%x: Tensor(fp16), %state: State(fp16)) -> {
            %read_state(fp16) = read_state(%state)
            %add(fp16) = add(%read_state, %x)
            %update(fp16) = coreml_update_state(%state, %add)
            %mul(fp16) = mul(%update, 5)
        } -> (%mul)
        """
        mlmodel = ct.convert(
            float32_buffer_model,
            inputs=[
                ct.TensorType(shape=(3,)),
            ],
            states=[
                ct.StateType(
                    wrapped_type=ct.TensorType(
                        shape=(3,),
                    ),
                    name="state",
                ),
            ],
            minimum_deployment_target=ct.target.iOS18,
            compute_units=compute_unit,
        )
        prog = mlmodel._mil_program
        test_valid_prog(prog)
        block = prog.functions["main"]
        assert block.inputs["x"].sym_type.get_primitive() == types.fp16
        assert (
            block.inputs["state_workaround"].sym_type.wrapped_type().get_primitive() == types.fp16
        )
        assert block.outputs[0].dtype == types.fp16
        verify_prediction(mlmodel)


        """
        fp16 state and fp32 input

        (%x: Tensor(fp32), %state: State(fp16)) -> {
            %read_state(fp16) = read_state(%state)
            %x_cast(fp16) = cast(%x)
            %add(fp16) = add(%read_state, %x_cast)
            %update(fp16) = coreml_update_state(%state, %add)
            %mul(fp16) = mul(%update, 5)
        } -> (%mul)
        """
        mlmodel = ct.convert(
            float32_buffer_model,
            inputs=[
                ct.TensorType(shape=(3,), dtype=np.float32),
            ],
            states=[
                ct.StateType(
                    wrapped_type=ct.TensorType(shape=(3,), dtype=np.float16), name="state"
                ),
            ],
            minimum_deployment_target=ct.target.iOS18,
            compute_units=compute_unit,
        )
        prog = mlmodel._mil_program
        expected_ops = [
            "read_state",
            "cast",
            "add",
            "coreml_update_state",
            "mul",
        ]
        test_valid_prog(prog, expected_ops)
        block = prog.functions["main"]
        assert block.inputs["x"].sym_type.get_primitive() == types.fp32
        assert (
            block.inputs["state_workaround"].sym_type.wrapped_type().get_primitive() == types.fp16
        )
        assert prog.find_ops("cast")[0].x.op is None
        assert block.outputs[0].dtype == types.fp16
        verify_prediction(mlmodel)

        """
        fp32 state and fp16 input. This is a rare corner case that shouldn't
        happend often.
        fp32 state is not supported through runtime.

        (%x: Tensor(fp16), %state: State(fp32)) -> {
            %read_state(fp32) = read_state(%state)
            %read_state_cast(fp16) = cast(read_state)
            %add(fp16) = add(%read_state_casr, %x)
            %add_cast(fp32) = cast(%add)
            %update(fp32) = coreml_update_state(%state, %add_cast)
            %update_cast(fp16) = cast(%update)
            %mul(fp16) = mul(%update_cast, 5)
        } -> (%mul)
        """
        prog = ct.convert(
            float32_buffer_model,
            inputs=[
                ct.TensorType(shape=(3,), dtype=np.float16),
            ],
            states=[
                ct.StateType(
                    wrapped_type=ct.TensorType(shape=(3,), dtype=np.float32), name="state"
                ),
            ],
            minimum_deployment_target=ct.target.iOS18,
            convert_to="milinternal",
        )
        expected_ops = [
            "read_state",
            "cast",
            "add",
            "cast",
            "coreml_update_state",
            "cast",
            "mul",
        ]
        test_valid_prog(prog, expected_ops)
        block = prog.functions["main"]
        assert block.inputs["x"].sym_type.get_primitive() == types.fp16
        assert (
            block.inputs["state_workaround"].sym_type.wrapped_type().get_primitive() == types.fp32
        )
        assert prog.find_ops("cast")[0].x.op.op_type == "read_state"
        assert prog.find_ops("cast")[1].x.op.op_type == "add"
        assert prog.find_ops("cast")[2].x.op.op_type == "coreml_update_state"
        assert block.outputs[0].dtype == types.fp16

        Tests for different combination of input dtypes.
        
coremltools/converters/mil/frontend/torch/test/test_torch_stateful_model.py