[BUG] dynamic_one_shot sample probability not always bounded by 1.

[dwierichs]

Expected behavior

dynamic_one_shot always computes valid MCM circuits successfully.

Actual behavior

For (very rare) configurations, the computed postselection probability can lie outside of $[0, 1]$ due to numerical rounding issues. This raises an error in the used binomial sampler.

Additional information

It was quite hard to find a point where this issue occurs. Deleting a few last digits from the weights array below already is enough to get to a successful evaluation.

Source code

import pennylane as qml
from pennylane import numpy as np

shots = 100
@qml.qnode(qml.device('default.qubit', shots=shots, seed=92552))
def qnode(weights):
    qml.BasicEntanglerLayers(weights, wires=range(2))
    mcms = [qml.measure(wire) for wire in range(2)]
    qml.cond(sum(mcms) == 2, qml.X)(wires=2)
    qml.measure(2, postselect=1)  
    return qml.expval(qml.Z(0))


weights = np.array([[0.24576389708811397, 0.364793497192904  ],
        [1.9671589829196052 , 0.3222075661288294 ],
        [0.7146512766410917 , 0.9736026815829427 ]], requires_grad=True)


qnode(weights)

Tracebacks

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[1], line 19
     11     return qml.expval(qml.Z(0))
     14 weights_autograd = np.array([[0.24576389708811397, 0.364793497192904  ],
     15         [1.9671589829196052 , 0.3222075661288294 ],
     16         [0.7146512766410917 , 0.9736026815829427 ]], requires_grad=True)
---> 19 qnode(weights_autograd)

File ~/repos/pennylane/pennylane/workflow/qnode.py:1164, 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 ~/repos/pennylane/pennylane/workflow/qnode.py:1150, 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 ~/repos/pennylane/pennylane/workflow/qnode.py:1103, 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 ~/repos/pennylane/pennylane/workflow/execution.py:835, 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 ~/repos/pennylane/pennylane/workflow/interfaces/autograd.py:147, 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 ~/venvs/dev/lib/python3.10/site-packages/autograd/tracer.py:48, in primitive.<locals>.f_wrapped(*args, **kwargs)
     46     return new_box(ans, trace, node)
     47 else:
---> 48     return f_raw(*args, **kwargs)

File ~/repos/pennylane/pennylane/workflow/interfaces/autograd.py:168, 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 ~/repos/pennylane/pennylane/workflow/execution.py:316, in _make_inner_execute.<locals>.inner_execute(tapes, **_)
    313     transformed_tapes = tuple(expand_fn(t) for t in transformed_tapes)
    315 if transformed_tapes:
--> 316     results = device_execution(transformed_tapes)
    317 else:
    318     results = ()

File ~/repos/pennylane/pennylane/devices/modifiers/simulator_tracking.py:30, in _track_execute.<locals>.execute(self, circuits, execution_config)
     28 @wraps(untracked_execute)
     29 def execute(self, circuits, execution_config=DefaultExecutionConfig):
---> 30     results = untracked_execute(self, circuits, execution_config)
     31     if isinstance(circuits, QuantumScript):
     32         batch = (circuits,)

File ~/repos/pennylane/pennylane/devices/modifiers/single_tape_support.py:32, in _make_execute.<locals>.execute(self, circuits, execution_config)
     30     is_single_circuit = True
     31     circuits = (circuits,)
---> 32 results = batch_execute(self, circuits, execution_config)
     33 return results[0] if is_single_circuit else results

File ~/repos/pennylane/pennylane/logging/decorators.py:61, in log_string_debug_func.<locals>.wrapper_entry(*args, **kwargs)
     54     s_caller = "::L".join(
     55         [str(i) for i in inspect.getouterframes(inspect.currentframe(), 2)[1][1:3]]
     56     )
     57     lgr.debug(
     58         f"Calling {f_string} from {s_caller}",
     59         **_debug_log_kwargs,
     60     )
---> 61 return func(*args, **kwargs)

File ~/repos/pennylane/pennylane/devices/default_qubit.py:597, in DefaultQubit.execute(self, circuits, execution_config)
    594 prng_keys = [self.get_prng_keys()[0] for _ in range(len(circuits))]
    596 if max_workers is None:
--> 597     return tuple(
    598         _simulate_wrapper(
    599             c,
    600             {
    601                 "rng": self._rng,
    602                 "debugger": self._debugger,
    603                 "interface": interface,
    604                 "state_cache": self._state_cache,
    605                 "prng_key": _key,
    606                 "mcm_method": execution_config.mcm_config.mcm_method,
    607                 "postselect_mode": execution_config.mcm_config.postselect_mode,
    608             },
    609         )
    610         for c, _key in zip(circuits, prng_keys)
    611     )
    613 vanilla_circuits = convert_to_numpy_parameters(circuits)[0]
    614 seeds = self._rng.integers(2**31 - 1, size=len(vanilla_circuits))

File ~/repos/pennylane/pennylane/devices/default_qubit.py:598, in <genexpr>(.0)
    594 prng_keys = [self.get_prng_keys()[0] for _ in range(len(circuits))]
    596 if max_workers is None:
    597     return tuple(
--> 598         _simulate_wrapper(
    599             c,
    600             {
    601                 "rng": self._rng,
    602                 "debugger": self._debugger,
    603                 "interface": interface,
    604                 "state_cache": self._state_cache,
    605                 "prng_key": _key,
    606                 "mcm_method": execution_config.mcm_config.mcm_method,
    607                 "postselect_mode": execution_config.mcm_config.postselect_mode,
    608             },
    609         )
    610         for c, _key in zip(circuits, prng_keys)
    611     )
    613 vanilla_circuits = convert_to_numpy_parameters(circuits)[0]
    614 seeds = self._rng.integers(2**31 - 1, size=len(vanilla_circuits))

File ~/repos/pennylane/pennylane/devices/default_qubit.py:863, in _simulate_wrapper(circuit, kwargs)
    862 def _simulate_wrapper(circuit, kwargs):
--> 863     return simulate(circuit, **kwargs)

File ~/repos/pennylane/pennylane/logging/decorators.py:61, in log_string_debug_func.<locals>.wrapper_entry(*args, **kwargs)
     54     s_caller = "::L".join(
     55         [str(i) for i in inspect.getouterframes(inspect.currentframe(), 2)[1][1:3]]
     56     )
     57     lgr.debug(
     58         f"Calling {f_string} from {s_caller}",
     59         **_debug_log_kwargs,
     60     )
---> 61 return func(*args, **kwargs)

File ~/repos/pennylane/pennylane/devices/qubit/simulate.py:347, in simulate(circuit, debugger, state_cache, **execution_kwargs)
    344     else:
    345         for i in range(circuit.shots.total_shots):
    346             results.append(
--> 347                 simulate_one_shot_native_mcm(
    348                     aux_circ, debugger=debugger, prng_key=keys[i], **execution_kwargs
    349                 )
    350             )
    351     return tuple(results)
    353 ops_key, meas_key = jax_random_split(prng_key)

File ~/repos/pennylane/pennylane/logging/decorators.py:61, in log_string_debug_func.<locals>.wrapper_entry(*args, **kwargs)
     54     s_caller = "::L".join(
     55         [str(i) for i in inspect.getouterframes(inspect.currentframe(), 2)[1][1:3]]
     56     )
     57     lgr.debug(
     58         f"Calling {f_string} from {s_caller}",
     59         **_debug_log_kwargs,
     60     )
---> 61 return func(*args, **kwargs)

File ~/repos/pennylane/pennylane/devices/qubit/simulate.py:753, in simulate_one_shot_native_mcm(circuit, debugger, **execution_kwargs)
    751 ops_key, meas_key = jax_random_split(prng_key)
    752 mid_measurements = {}
--> 753 state, is_state_batched = get_final_state(
    754     circuit,
    755     debugger=debugger,
    756     mid_measurements=mid_measurements,
    757     prng_key=ops_key,
    758     **execution_kwargs,
    759 )
    760 return measure_final_state(
    761     circuit,
    762     state,
   (...)
    766     **execution_kwargs,
    767 )

File ~/repos/pennylane/pennylane/logging/decorators.py:61, in log_string_debug_func.<locals>.wrapper_entry(*args, **kwargs)
     54     s_caller = "::L".join(
     55         [str(i) for i in inspect.getouterframes(inspect.currentframe(), 2)[1][1:3]]
     56     )
     57     lgr.debug(
     58         f"Calling {f_string} from {s_caller}",
     59         **_debug_log_kwargs,
     60     )
---> 61 return func(*args, **kwargs)

File ~/repos/pennylane/pennylane/devices/qubit/simulate.py:174, in get_final_state(circuit, debugger, **execution_kwargs)
    172 if isinstance(op, MidMeasureMP):
    173     prng_key, key = jax_random_split(prng_key)
--> 174 state = apply_operation(
    175     op,
    176     state,
    177     is_state_batched=is_state_batched,
    178     debugger=debugger,
    179     prng_key=key,
    180     tape_shots=circuit.shots,
    181     **execution_kwargs,
    182 )
    183 # Handle postselection on mid-circuit measurements
    184 if isinstance(op, qml.Projector):

File /usr/lib/python3.10/functools.py:889, in singledispatch.<locals>.wrapper(*args, **kw)
    885 if not args:
    886     raise TypeError(f'{funcname} requires at least '
    887                     '1 positional argument')
--> 889 return dispatch(args[0].__class__)(*args, **kw)

File ~/repos/pennylane/pennylane/devices/qubit/apply_operation.py:345, in apply_mid_measure(op, state, is_state_batched, debugger, **execution_kwargs)
    343     else:
    344         binomial_fn = np.random.binomial if rng is None else rng.binomial
--> 345     sample = binomial_fn(1, 1 - prob0)
    346 mid_measurements[op] = sample
    348 # Using apply_operation(qml.QubitUnitary,...) instead of apply_operation(qml.Projector([sample], wire),...)
    349 # to select the sample branch enables jax.jit and prevents it from using Python callbacks

File numpy/random/_generator.pyx:3020, in numpy.random._generator.Generator.binomial()

File _common.pyx:430, in numpy.random._common.check_constraint()

ValueError: p < 0, p > 1 or p is NaN

System information

pl dev/master

Existing GitHub issues

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