Repository: PennyLaneAI/pennylane

Issue:
`default.qubit`  returns analytic probabilities when `analytic=False`
Currently, when `qml.probs()` is returned from a QNode, the `QubitDevice.probability()` abstract method is returned directly to the user.

For hardware-like devices that don't support analytic mode, this should behave as expected; the probability will be computed on the device by binning the generated computational basis samples.

However, for _analytic_ devices, such as `default.qubit`, care must be taken here. Currently, there is a bug in `DefaultQubit.probability()`, in that the analytic probability is _always_ returned, regardless what the value of `device.analytic` is:

```python
def probability(self, wires=None):
    if self._state is None:
        return None

    wires = wires or range(self.num_wires)
    prob = self.marginal_prob(np.abs(self._state) ** 2, wires)
    return prob
```

This should be modified as follows:
```python
def probability(self, wires=None):
    if self._state is None:
        return None

    wires = wires or range(self.num_wires)

    if self.analytic:
        prob = self.marginal_prob(np.abs(self._state) ** 2, wires)
        return prob

    # non-analytic mode, estimate the probability from the generated samples

    # consider only the requested wires
    wires = np.hstack(wires)
    samples = self._samples[:, np.array(wires)]

    # convert samples from a list of 0, 1 integers,
    # to base 10 representation
    unraveled_indices = [2] * len(wires)
    indices = np.ravel_multi_index(samples.T, unraveled_indices)

    # count the basis state occurrences, and construct
    # the probability vector
    basis_states, counts = np.unique(indices, return_counts=True)
    prob = np.zeros([len(wires)**2], dtype=np.float64)
    prob[basis_states] = counts/self.shots
    return prob
```


Discussion:

Resolving patch:
diff --git a/pennylane/_qubit_device.py b/pennylane/_qubit_device.py
--- a/pennylane/_qubit_device.py
+++ b/pennylane/_qubit_device.py
@@ -280,7 +280,15 @@ def generate_samples(self):
              array[complex]: array of samples in the shape ``(dev.shots, dev.num_wires)``
         """
         number_of_states = 2 ** self.num_wires
-        rotated_prob = self.probability()
+
+        # Check if device has `analytic_probability` method, otherwise call `probability`
+        # TODO: try-except statement should be removed when enforcing an `analytic_probabily`
+        # method for all analytic devices => should always call `self.analytic_probability`
+        try:
+            rotated_prob = self.analytic_probability()
+        except AttributeError:
+            rotated_prob = self.probability()
+
         samples = self.sample_basis_states(number_of_states, rotated_prob)
         return QubitDevice.states_to_binary(samples, self.num_wires)
 
diff --git a/pennylane/plugins/default_qubit.py b/pennylane/plugins/default_qubit.py
--- a/pennylane/plugins/default_qubit.py
+++ b/pennylane/plugins/default_qubit.py
@@ -240,10 +240,38 @@ def reset(self):
         self._state[0] = 1
         self._pre_rotated_state = self._state
 
-    def probability(self, wires=None):
+    def analytic_probability(self, wires=None):
+        """Return the (marginal) analytic probability of each computational basis state."""
         if self._state is None:
             return None
 
         wires = wires or range(self.num_wires)
+
         prob = self.marginal_prob(np.abs(self._state) ** 2, wires)
         return prob
+
+    def estimate_probability(self, wires=None):
+        """Return the estimated probability of each computational basis state using the generated samples."""
+
+        # consider only the requested wires
+        wires = np.hstack(wires)
+
+        samples = self._samples[:, np.array(wires)]
+
+        # convert samples from a list of 0, 1 integers, to base 10 representation
+        unraveled_indices = [2] * len(wires)
+        indices = np.ravel_multi_index(samples.T, unraveled_indices)
+
+        # count the basis state occurrences, and construct the probability vector
+        basis_states, counts = np.unique(indices, return_counts=True)
+        prob = np.zeros([len(wires) ** 2], dtype=np.float64)
+        prob[basis_states] = counts / self.shots
+        return prob
+
+    def probability(self, wires=None):
+        wires = wires or range(self.num_wires)
+
+        if self.analytic:
+            return self.analytic_probability(wires=wires)
+
+        return self.estimate_probability(wires=wires)


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -172,6 +172,17 @@ def mock_device(monkeypatch):
         m.setattr(dev, 'capabilities', lambda cls: {"model": "qubit"})
         yield qml.Device(wires=2)
 
+@pytest.fixture(scope="function")
+def mock_qubit_device(monkeypatch):
+    """A mock instance of the abstract QubitDevice class"""
+
+    with monkeypatch.context() as m:
+        dev = qml.QubitDevice
+        m.setattr(dev, '__abstractmethods__', frozenset())
+        m.setattr(dev, 'short_name', 'mock_device')
+        m.setattr(dev, 'capabilities', lambda cls: {"model": "qubit"})
+        yield qml.QubitDevice(wires=2)
+
 @pytest.fixture
 def tear_down_hermitian():
     yield None
diff --git a/tests/test_default_qubit.py b/tests/test_default_qubit.py
--- a/tests/test_default_qubit.py
+++ b/tests/test_default_qubit.py
@@ -1612,3 +1612,61 @@ def test_hermitian(self, theta, phi, varphi, monkeypatch, tol):
             )
         ) / 16
         assert np.allclose(var, expected, atol=tol, rtol=0)
+
+class TestProbabilityIntegration:
+    """Test probability method for when analytic is True/False"""
+
+    def mock_analytic_counter(self, wires=None):
+        self.analytic_counter += 1
+        return np.array([1, 0, 0, 0], dtype=float)
+
+    @pytest.mark.parametrize("x", [[0.2, 0.5], [0.4, 0.9], [0.8, 0.3]])
+    def test_probability(self, x, tol):
+        """Test that the probability function works when analytic=False"""
+        dev = qml.device("default.qubit", wires=2, analytic=False)
+        dev_analytic = qml.device("default.qubit", wires=2, analytic=True)
+
+        def circuit(x):
+            qml.RX(x[0], wires=0)
+            qml.RY(x[1], wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.probs(wires=[0, 1])
+
+        prob = qml.QNode(circuit, dev)
+        prob_analytic = qml.QNode(circuit, dev_analytic)
+
+        assert np.isclose(prob(x).sum(), 1, atol=tol, rtol=0)
+        assert np.allclose(prob_analytic(x), prob(x), atol=0.1, rtol=0)
+        assert not np.array_equal(prob_analytic(x), prob(x))
+
+    @pytest.mark.parametrize("analytic", [True, False])
+    def test_call_generate_samples(self, analytic, monkeypatch):
+        """Test analytic_probability call when generating samples"""
+        self.analytic_counter = False
+
+        dev = qml.device("default.qubit", wires=2, analytic=analytic)
+        monkeypatch.setattr(dev, "analytic_probability", self.mock_analytic_counter)
+
+        # generate samples through `generate_samples` (using 'analytic_probability')
+        dev.generate_samples()
+
+        # should call `analytic_probability` once through `generate_samples`
+        assert self.analytic_counter == 1
+
+    # TODO: remove test when enforcing `analytic_probability` method for all analytic devices
+    @pytest.mark.parametrize("analytic", [True, False])
+    def test_call_generate_samples_no_analytic_fn(self, analytic, mock_qubit_device, monkeypatch):
+        """Test analytic_probability when device has no method called analytic_probability"""
+        self.analytic_counter = 0
+
+        dev = qml.device("default.qubit", wires=2, analytic=analytic)
+        monkeypatch.setattr(dev, "analytic_probability", self.mock_analytic_counter)
+
+        # generate samples from a qubit device with no 'analytic_probabiliy' method
+        dev._samples = mock_qubit_device.generate_samples()
+        dev.probability()
+
+        # should call `mock_analytic_probability` through `probability` once only if analytic=True
+        # and never through `generate_samples` since the `mock_qubit_device` has no mocked
+        # `analytic_probabiliy` method.
+        assert self.analytic_counter == analytic