Unexpected output from dynamic allocation with sample measurement process

[sengthai]

qml.allocate registers new auxiliary qubits, which are then deallocated after use.
In this circuit, when we apply an operation to an auxiliary qubit, it seems affect the original qubits.
It is supposed to return [0,0] ten times.

We expect this error to be caused by the runtime.

import qml.capture.enable()
    
@qml.qjit
@qml.qnode(qml.device("lightning.qubit", wires=3, shots=10))
def circuit():
    with qml.allocate(2) as qs:
        qml.X(qs[1])

    return qml.sample(wires=[0, 1])

Output:

[[0 0]
 [0 1]
 [0 0]
 [1 0]
 [0 0]
 [0 0]
 [0 1]
 [0 0]
 [1 0]
 [0 0]]

[rniczh]

It probably has some bug in the GenerateSample(), counts also encounter the same issue and it also use the GenerateSample().

void LightningSimulator::PartialSample(DataView<double, 2> &samples,
                                       const std::vector<QubitIdType> &wires) {
    const size_t numWires = wires.size();
    const size_t numQubits = this->GetNumQubits();

    RT_FAIL_IF(numWires > numQubits, "Invalid number of wires");
    RT_FAIL_IF(!isValidQubits(wires), "Invalid given wires to measure");
    RT_FAIL_IF(samples.size() != device_shots * numWires,
               "Invalid size for the pre-allocated partial-samples");

    // get device wires
    auto &&dev_wires = getDeviceWires(wires);

    auto li_samples = this->GenerateSamples(device_shots);

    // The lightning samples are layed out as a single vector of size
    // shots*qubits, where each element represents a single bit. The
    // corresponding shape is (shots, qubits). Gather the desired bits
    // corresponding to the input wires into a bitstring.
    auto samplesIter = samples.begin();
    for (size_t shot = 0; shot < device_shots; shot++) {
        for (auto wire : dev_wires) {
            *(samplesIter++) =
                static_cast<double>(li_samples[shot * numQubits + wire]);
        }
    }
}

Turns out if we do the partialProbs first and base on the result to generate samples can solve this issue:
In RuntimeCAPI.cpp. I opened a PR but I think that should be resolved from lightning side

size_t shots = getQuantumDevicePtr()->GetDeviceShots();
size_t numWires = wires.size();
size_t numBasisStates = 1UL << numWires;

// Get marginal probabilities for wires
std::vector<double> probs(numBasisStates);
size_t probs_sizes[1] = {numBasisStates};
size_t probs_strides[1] = {1};
DataView<double, 1> probs_view(probs.data(), 0, probs_sizes, probs_strides);
getQuantumDevicePtr()->PartialProbs(probs_view, wires);

// Sample from marginal probability distribution
std::random_device rd;
std::mt19937 gen(rd());
std::discrete_distribution<size_t> dist(probs.begin(), probs.end());

// Generate samples
auto samplesIter = view.begin();
for (size_t shot = 0; shot < shots; shot++) {
    size_t basis_state = dist(gen);
    for (size_t wire_idx = 0; wire_idx < numWires; wire_idx++) {
        *(samplesIter++) = static_cast<double>((basis_state >> wire_idx) & 1UL);
    }
}

EDIT:
We opened the PR in lightning to address this issue by adjusting state vector instead. PennyLaneAI/pennylane-lightning#1321