[BUG] dynamic_one_shot is ignored when applied as a QNode transform.

[astralcai]

Expected behavior

dynamic_one_shot works as a QNode transform.

Actual behavior

@partial(qml.dynamic_one_shot, postselect_mode="fill-shots")
@qml.qnode(qml.device("default.qubit", shots=100))
def circuit():
    qml.H(0)
    _ = qml.measure(0, postselect=1)
    return qml.sample(wires=[0])

>>> res = circuit()
>>> res.shape # expecting (100,) because "fill-shots"
(53,)

Additional information

This is due to default.qubit automatically applying its own mid_circuit_measurements transform in preprocess, which is prioritized over any user-applied dynamic_one_shot transform:

pennylane/pennylane/workflow/qnode.py

Line 924 in 147fb99

_prune_dynamic_transform(full_transform_program, inner_transform_program)

I would argue that it should be the other way around. If the user explicitly applies a dynamic_one_shot transform on the QNode, it should be used instead of the transform that already exists in the preprocess transform program.

Source code

No response

Tracebacks

No response

System information

Name: PennyLane
Version: 0.40.0.dev7
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: /Users/astral.cai/Workspace/pennylane/venv/lib/python3.10/site-packages
Editable project location: /Users/astral.cai/Workspace/pennylane
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, packaging, pennylane-lightning, requests, rustworkx, scipy, toml, typing-extensions
Required-by: PennyLane-Catalyst, PennyLane_Lightning, PennyLane_Lightning_Kokkos

Platform info:           macOS-15.1-arm64-arm-64bit
Python version:          3.10.14
Numpy version:           1.26.4
Scipy version:           1.12.0
Installed devices:
- default.clifford (PennyLane-0.40.0.dev8)
- default.gaussian (PennyLane-0.40.0.dev8)
- default.mixed (PennyLane-0.40.0.dev8)
- default.qubit (PennyLane-0.40.0.dev8)
- default.qutrit (PennyLane-0.40.0.dev8)
- default.qutrit.mixed (PennyLane-0.40.0.dev8)
- default.tensor (PennyLane-0.40.0.dev8)
- null.qubit (PennyLane-0.40.0.dev8)
- reference.qubit (PennyLane-0.40.0.dev8)
- lightning.qubit (PennyLane_Lightning-0.39.0.dev48)
- nvidia.custatevec (PennyLane-Catalyst-0.9.0.dev41)
- nvidia.cutensornet (PennyLane-Catalyst-0.9.0.dev41)
- oqc.cloud (PennyLane-Catalyst-0.9.0.dev41)
- softwareq.qpp (PennyLane-Catalyst-0.9.0.dev41)
- lightning.kokkos (PennyLane_Lightning_Kokkos-0.39.0.dev38)

Existing GitHub issues

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

[albi3ro]

Should we just forbid dynamic_one_shot from being used as a user transform?