[Feature]: Support multiple arrays in `state_output`

[CalMacCQ]

Problem Statement

If I understand correctly, its only possible to pass a single array of qubits to a state_output call. This limits the user's ability to obtain a statevector printout over multiple "registers" (aka arrays).

I think it should be possible to pass selene the appropriate indices.

An example of how I'm working around this currently:

from guppylang import guppy
from guppylang.std.array import array
from guppylang.std.quantum import qubit, discard_array, h, cx
from guppylang.std.debug import state_output

@guppy
def main() -> None:
    control_block = array(qubit() for _ in range(4))
    target_block = array(qubit() for _ in range(4))
    for i in range(4):
        h(control_block[i])
        cx(control_block[i], target_block[i])

        
    # Each state_result call is applied to a single array of qubits.
    state_output("control", control_block)
    state_output("target", target_block)
    state_output("total", target_block) # hack

    discard_array(control_block)
    discard_array(target_block)

instance = build(main.compile())
seeded_stim_instance = Stim(random_seed=seed)
output = instance.run(simulator=seeded_stim_instance, n_qubits=8)
states_dict = seeded_stim_instance.extract_states_dict(output)

I can then hack SeleneStimState.specified_qubits as follows

total = states_dict["total"]
control_qubits = states_dict["control"].specified_qubits
target_qubits = states_dict["target"].specified_qubits
total.specified_qubits = control_qubits + target_qubits

Of course it would be nicer to do

@guppy
def main() -> None:
    control_block = array(qubit() for _ in range(4))
    target_block = array(qubit() for _ in range(4))
    for i in range(4):
        h(control_block[i])
        cx(control_block[i], target_block[i])

    state_output("state", control_block, target_block) # unsupported currently

    discard_array(control_block)
    discard_array(target_block)

instance = build(main.compile())
seeded_stim_instance = Stim(random_seed=seed)
output = instance.run(simulator=seeded_stim_instance, n_qubits=8)
states_dict = seeded_stim_instance.extract_states_dict(output)
state = states_dict["state"] 

In the later snippet we have a single state_output call rather than three. We also don't have to manually edit the SeleneStimState.specified_qubits attribute.

Proposed Solution

No response

Component

No response

Effort Estimate

None

Additional Context

No response