[BUG] Array indexing error when sampling after mid-circuit measurements are performed on other wires

[glassnotes]

Expected behavior

When performing a circuit with mid-circuit measurements, returning sample measurements on which no mid-circuit measurements have been performed should work.

Actual behavior

An error (which appears related to the mid-circuit measurement collection) is thrown.

Additional information

No response

Source code

dev = qml.device("default.qubit", wires=2, shots=1)

# This version does not work; nor does sampling on wire 0, even
# if we set reset=True in qml.measure
@qml.qnode(dev)
def apply_mcm():
    qml.Hadamard(wires=0)
    m0 = qml.measure(wires=0)
    return qml.sample(wires=1)

# Using probs instead of samples does work
@qml.qnode(dev)
def apply_mcm():
    qml.Hadamard(wires=0)
    m0 = qml.measure(wires=0)
    return qml.probs(wires=1)

# Removing the MCM works with samples
@qml.qnode(dev)
def apply_mcm():
    qml.Hadamard(wires=0)
    return qml.sample(wires=1)

# Sampling the MCM itself works
@qml.qnode(dev)
def apply_mcm():
    qml.Hadamard(wires=0)
    m0 = qml.measure(wires=0)
    return qml.sample(m0)

Tracebacks

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
Cell In[53], line 1
----> 1 apply_mcm()

File [~/.conda/envs/qecc/lib/python3.12/site-packages/pennylane/workflow/qnode.py:1164](http://localhost:8888/home/olivia/.conda/envs/qecc/lib/python3.12/site-packages/pennylane/workflow/qnode.py#line=1163), in QNode.__call__(self, *args, **kwargs)
   1162 if qml.capture.enabled():
   1163     return qml.capture.qnode_call(self, *args, **kwargs)
-> 1164 return self._impl_call(*args, **kwargs)

File [~/.conda/envs/qecc/lib/python3.12/site-packages/pennylane/workflow/qnode.py:1150](http://localhost:8888/home/olivia/.conda/envs/qecc/lib/python3.12/site-packages/pennylane/workflow/qnode.py#line=1149), in QNode._impl_call(self, *args, **kwargs)
   1147 self._update_gradient_fn(shots=override_shots, tape=self._tape)
   1149 try:
-> 1150     res = self._execution_component(args, kwargs, override_shots=override_shots)
   1151 finally:
   1152     if old_interface == "auto":

File [~/.conda/envs/qecc/lib/python3.12/site-packages/pennylane/workflow/qnode.py:1103](http://localhost:8888/home/olivia/.conda/envs/qecc/lib/python3.12/site-packages/pennylane/workflow/qnode.py#line=1102), in QNode._execution_component(self, args, kwargs, override_shots)
   1100 _prune_dynamic_transform(full_transform_program, inner_transform_program)
   1102 # pylint: disable=unexpected-keyword-arg
-> 1103 res = qml.execute(
   1104     (self._tape,),
   1105     device=self.device,
   1106     gradient_fn=self.gradient_fn,
   1107     interface=self.interface,
   1108     transform_program=full_transform_program,
   1109     inner_transform=inner_transform_program,
   1110     config=config,
   1111     gradient_kwargs=self.gradient_kwargs,
   1112     override_shots=override_shots,
   1113     **self.execute_kwargs,
   1114 )
   1115 res = res[0]
   1117 # convert result to the interface in case the qfunc has no parameters

File [~/.conda/envs/qecc/lib/python3.12/site-packages/pennylane/workflow/execution.py:835](http://localhost:8888/home/olivia/.conda/envs/qecc/lib/python3.12/site-packages/pennylane/workflow/execution.py#line=834), in execute(tapes, device, gradient_fn, interface, transform_program, inner_transform, config, grad_on_execution, gradient_kwargs, cache, cachesize, max_diff, override_shots, expand_fn, max_expansion, device_batch_transform, device_vjp, mcm_config)
    827 ml_boundary_execute = _get_ml_boundary_execute(
    828     interface,
    829     _grad_on_execution,
    830     config.use_device_jacobian_product,
    831     differentiable=max_diff > 1,
    832 )
    834 if interface in jpc_interfaces:
--> 835     results = ml_boundary_execute(tapes, execute_fn, jpc, device=device)
    836 else:
    837     results = ml_boundary_execute(
    838         tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1, max_diff=max_diff
    839     )

File [~/.conda/envs/qecc/lib/python3.12/site-packages/pennylane/workflow/interfaces/autograd.py:147](http://localhost:8888/home/olivia/.conda/envs/qecc/lib/python3.12/site-packages/pennylane/workflow/interfaces/autograd.py#line=146), in autograd_execute(tapes, execute_fn, jpc, device)
    142 # pylint misidentifies autograd.builtins as a dict
    143 # pylint: disable=no-member
    144 parameters = autograd.builtins.tuple(
    145     [autograd.builtins.list(t.get_parameters()) for t in tapes]
    146 )
--> 147 return _execute(parameters, tuple(tapes), execute_fn, jpc)

File [~/.conda/envs/qecc/lib/python3.12/site-packages/autograd/tracer.py:48](http://localhost:8888/home/olivia/.conda/envs/qecc/lib/python3.12/site-packages/autograd/tracer.py#line=47), in primitive.<locals>.f_wrapped(*args, **kwargs)
     46     return new_box(ans, trace, node)
     47 else:
---> 48     return f_raw(*args, **kwargs)

File [~/.conda/envs/qecc/lib/python3.12/site-packages/pennylane/workflow/interfaces/autograd.py:168](http://localhost:8888/home/olivia/.conda/envs/qecc/lib/python3.12/site-packages/pennylane/workflow/interfaces/autograd.py#line=167), in _execute(parameters, tapes, execute_fn, jpc)
    150 @autograd.extend.primitive
    151 def _execute(
    152     parameters,
   (...)
    155     jpc,
    156 ):  # pylint: disable=unused-argument
    157     """Autodifferentiable wrapper around a way of executing tapes.
    158 
    159     Args:
   (...)
    166 
    167     """
--> 168     return execute_fn(tapes)

File [~/.conda/envs/qecc/lib/python3.12/site-packages/pennylane/workflow/execution.py:320](http://localhost:8888/home/olivia/.conda/envs/qecc/lib/python3.12/site-packages/pennylane/workflow/execution.py#line=319), in _make_inner_execute.<locals>.inner_execute(tapes, **_)
    317 else:
    318     results = ()
--> 320 return transform_post_processing(results)

File [~/.conda/envs/qecc/lib/python3.12/site-packages/pennylane/transforms/core/transform_program.py:86](http://localhost:8888/home/olivia/.conda/envs/qecc/lib/python3.12/site-packages/pennylane/transforms/core/transform_program.py#line=85), in _apply_postprocessing_stack(results, postprocessing_stack)
     63 """Applies the postprocessing and cotransform postprocessing functions in a Last-In-First-Out LIFO manner.
     64 
     65 Args:
   (...)
     83 
     84 """
     85 for postprocessing in reversed(postprocessing_stack):
---> 86     results = postprocessing(results)
     87 return results

File [~/.conda/envs/qecc/lib/python3.12/site-packages/pennylane/transforms/core/transform_program.py:56](http://localhost:8888/home/olivia/.conda/envs/qecc/lib/python3.12/site-packages/pennylane/transforms/core/transform_program.py#line=55), in _batch_postprocessing(results, individual_fns, slices)
     30 def _batch_postprocessing(
     31     results: ResultBatch, individual_fns: List[PostProcessingFn], slices: List[slice]
     32 ) -> ResultBatch:
     33     """Broadcast individual post processing functions onto their respective tapes.
     34 
     35     Args:
   (...)
     54 
     55     """
---> 56     return tuple(fn(results[sl]) for fn, sl in zip(individual_fns, slices))

File [~/.conda/envs/qecc/lib/python3.12/site-packages/pennylane/transforms/core/transform_program.py:56](http://localhost:8888/home/olivia/.conda/envs/qecc/lib/python3.12/site-packages/pennylane/transforms/core/transform_program.py#line=55), in <genexpr>(.0)
     30 def _batch_postprocessing(
     31     results: ResultBatch, individual_fns: List[PostProcessingFn], slices: List[slice]
     32 ) -> ResultBatch:
     33     """Broadcast individual post processing functions onto their respective tapes.
     34 
     35     Args:
   (...)
     54 
     55     """
---> 56     return tuple(fn(results[sl]) for fn, sl in zip(individual_fns, slices))

File [~/.conda/envs/qecc/lib/python3.12/site-packages/pennylane/transforms/dynamic_one_shot.py:163](http://localhost:8888/home/olivia/.conda/envs/qecc/lib/python3.12/site-packages/pennylane/transforms/dynamic_one_shot.py#line=162), in dynamic_one_shot.<locals>.processing_fn(results, has_partitioned_shots, batched_results)
    159 else:
    160     results = [
    161         reshape_data(tuple(res[i] for res in results)) for i, _ in enumerate(results[0])
    162     ]
--> 163 return parse_native_mid_circuit_measurements(tape, aux_tapes, results, interface=interface)

File [~/.conda/envs/qecc/lib/python3.12/site-packages/pennylane/transforms/dynamic_one_shot.py:303](http://localhost:8888/home/olivia/.conda/envs/qecc/lib/python3.12/site-packages/pennylane/transforms/dynamic_one_shot.py#line=302), in parse_native_mid_circuit_measurements(circuit, aux_tapes, results, interface)
    301             continue
    302         result = qml.math.squeeze(result)
--> 303     meas = gather_non_mcm(m, result, is_valid)
    304     m_count += 1
    305 if isinstance(m, SampleMP):

File [~/.conda/envs/qecc/lib/python3.12/site-packages/pennylane/transforms/dynamic_one_shot.py:378](http://localhost:8888/home/olivia/.conda/envs/qecc/lib/python3.12/site-packages/pennylane/transforms/dynamic_one_shot.py#line=377), in gather_non_mcm(measurement, samples, is_valid)
    373     if is_interface_jax and samples.ndim == 2:
    374         is_valid = is_valid.reshape((-1, 1))
    375     return (
    376         qml.math.where(is_valid, samples, fill_in_value)
    377         if is_interface_jax
--> 378         else samples[is_valid]
    379     )
    380 # VarianceMP
    381 expval = qml.math.sum(samples * is_valid) [/](http://localhost:8888/) qml.math.sum(is_valid)

File [~/.conda/envs/qecc/lib/python3.12/site-packages/pennylane/numpy/tensor.py:187](http://localhost:8888/home/olivia/.conda/envs/qecc/lib/python3.12/site-packages/pennylane/numpy/tensor.py#line=186), in tensor.__getitem__(self, *args, **kwargs)
    186 def __getitem__(self, *args, **kwargs):
--> 187     item = super().__getitem__(*args, **kwargs)
    189     if not isinstance(item, tensor):
    190         item = tensor(item, requires_grad=self.requires_grad)

IndexError: too many indices for array: array is 0-dimensional, but 1 were indexed

System information

Name: PennyLane
Version: 0.37.0
Summary: PennyLane is a cross-platform Python library for quantum computing, quantum machine learning, and quantum chemistry. Train a quantum computer the same way as a neural network.
Home-page: https://github.com/PennyLaneAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: /home/olivia/.conda/envs/qecc/lib/python3.12/site-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, packaging, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing-extensions
Required-by: PennyLane_Lightning

Platform info:           Linux-6.8.0-38-generic-x86_64-with-glibc2.39
Python version:          3.12.4
Numpy version:           1.26.4
Scipy version:           1.14.0
Installed devices:
- lightning.qubit (PennyLane_Lightning-0.37.0)
- default.clifford (PennyLane-0.37.0)
- default.gaussian (PennyLane-0.37.0)
- default.mixed (PennyLane-0.37.0)
- default.qubit (PennyLane-0.37.0)
- default.qubit.autograd (PennyLane-0.37.0)
- default.qubit.jax (PennyLane-0.37.0)
- default.qubit.legacy (PennyLane-0.37.0)
- default.qubit.tf (PennyLane-0.37.0)
- default.qubit.torch (PennyLane-0.37.0)
- default.qutrit (PennyLane-0.37.0)
- default.qutrit.mixed (PennyLane-0.37.0)
- default.tensor (PennyLane-0.37.0)
- null.qubit (PennyLane-0.37.0)

Existing GitHub issues

• I have searched existing GitHub issues to make sure the issue does not already exist.

[albi3ro]

Given it works with two shots, I have a feeling its yet another squeezing issue.

[trbromley]

Thanks @glassnotes! How urgent / blocking is this for you?

Note that the problem seems specific to the new one-shot approach because the following works:

import pennylane as qml

dev = qml.device("default.qubit", wires=2, shots=1)

@qml.qnode(dev, mcm_method="deferred")
def apply_mcm():
    qml.Hadamard(wires=0)
    m0 = qml.measure(wires=0)
    return qml.sample(wires=1)

apply_mcm()

[glassnotes]

In our application we're working with quite a large circuit (>25 qubits), so deferred measurement isn't an option. But as @albi3ro notes, it works for multiple shots, so it's not an outright blocker any more (thanks for pointing this out!).