[BUG] `QubitDevice`/`QutritDevice` `probability(wires=0)` silently returns the full distribution instead of the marginal

[YujinSong-hep]

Expected behavior

The probability method documents that wires accepts a single wire label:

wires (Iterable[Number, str], Number, str, Wires): wires to return marginal probabilities for.
Wires not provided are traced out of the system.

So probability(wires=0) should return the marginal distribution for wire 0 only. For a
2-qutrit device this is a 3-element array; for a 2-qubit device a 2-element array.

Actual behavior

probability(wires=0) returns the full joint distribution over all wires, identical to
probability(wires=None). The user's wire specification is silently ignored:

probability(wires=None) length: 9
probability(wires=0)    length: 9   (expected 3 for one qutrit)
probability(wires=0) == full distribution: True

The integer wire label 0 is treated as "no wires given", so the call falls back to all device
wires. This affects any code that requests the marginal of wire 0 — the most common first wire.

Additional information

Root cause is a truthiness-based default check. In pennylane/devices/_qubit_device.py
(QubitDevice.probability, ~line 1270) and the equivalent in
pennylane/devices/_qutrit_device.py (QutritDevice.probability, ~line 265):

def probability(self, wires=None, shot_range=None, bin_size=None):
    wires = wires or self.wires   # BUG: integer wire label 0 is falsy -> replaced by self.wires
    ...

wires or self.wires cannot distinguish "the user passed wire 0" from "the user passed no
wires", because 0 is falsy in Python. The same pattern also appears in
QubitDevice.estimate_probability (~line 1179). The fix is an explicit None check:

wires = wires if wires is not None else self.wires

Source code

import numpy as np
import pennylane as qml

dev = qml.device('default.qutrit', wires=2)

@qml.qnode(dev)
def circuit():
    qml.TShift(0)               
    return qml.probs(wires=[0, 1])

circuit()   

full = np.asarray(dev.probability(wires=None))
w0   = np.asarray(dev.probability(wires=0))

print("probability(wires=None) length:", full.shape[0])
print("probability(wires=0)    length:", w0.shape[0], "(expected 3 for one qutrit)")
print("probability(wires=0) == full distribution:", np.allclose(w0, full))

Tracebacks

probability(wires=None) length: 9
probability(wires=0)    length: 9 (expected 3 for one qutrit)
probability(wires=0) == full distribution: True

System information

Version: 0.45.0
Platform info: macOS
Python version: 3.13.13

Existing GitHub issues

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

[CatalinaAlbornoz]

Thanks for reporting this bug @YujinSong-hep .