Repository: PennyLaneAI/pennylane

Issue:
[BUG] QNodes that return `qml.sample(obs)` and use `parameter-shift` are not differentiable 
### Expected behavior

A QNode with finite shots that returns `qml.sample(obs)` and uses the PR-rule should be differentiable.

### Actual behavior

It's not in general.

### Additional information

- Related issue: https://github.com/PennyLaneAI/pennylane/issues/2556
- Related forum post: https://discuss.pennylane.ai/t/error-unhashable-type-listandwires-must-be-hashable-got-0-1/3653/6

`qml.sample(qml.PauliZ(0))` returns a list of `int`s. Should it return a list of `float`s? The documentation for `qml.sample` seems to indicate that it should. This also happens for `default.qubit.legacy`

![Screenshot 2023-11-14 at 10 27 01 AM](https://github.com/PennyLaneAI/pennylane/assets/34751083/2b4d3712-59d6-4bc0-9c00-d4b55ce16d1e)

### Source code

```shell
dev = qml.device("default.qubit", shots=10)

@qml.qnode(dev, diff_method="parameter-shift")
def circuit(angle):
    qml.RX(angle, wires=0)
    return qml.sample(qml.PauliX(0))

def cost(angle):
    return np.mean(circuit(angle))

angle = np.array(0.1)

opt = qml.GradientDescentOptimizer(stepsize=0.5)
print(opt.step_and_cost(cost, angle))
```


### Tracebacks

```shell
---------------------------------------------------------------------------
UFuncTypeError                            Traceback (most recent call last)
/Users/Documents/notebook.ipynb Cell 2 line 1
     11 angle = np.array(0.1)
     13 opt = qml.GradientDescentOptimizer(stepsize=0.5)
---> 14 print(opt.step_and_cost(cost, angle))

File ~/.virtualenvs/pennylane-stable/lib/python3.9/site-packages/pennylane/optimize/gradient_descent.py:59, in GradientDescentOptimizer.step_and_cost(self, objective_fn, grad_fn, *args, **kwargs)
     39 def step_and_cost(self, objective_fn, *args, grad_fn=None, **kwargs):
     40     """Update trainable arguments with one step of the optimizer and return the corresponding
     41     objective function value prior to the step.
     42 
   (...)
     56         If single arg is provided, list [array] is replaced by array.
     57     """
---> 59     g, forward = self.compute_grad(objective_fn, args, kwargs, grad_fn=grad_fn)
     60     new_args = self.apply_grad(g, args)
     62     if forward is None:

File ~/.virtualenvs/pennylane-stable/lib/python3.9/site-packages/pennylane/optimize/gradient_descent.py:117, in GradientDescentOptimizer.compute_grad(objective_fn, args, kwargs, grad_fn)
     99 r"""Compute gradient of the objective function at the given point and return it along with
    100 the objective function forward pass (if available).
    101 
   (...)
    114     will not be evaluted and instead ``None`` will be returned.
...

File ~/.virtualenvs/pennylane-stable/lib/python3.9/site-packages/autograd/core.py:213, in VSpace._mut_add(self, x, y)
--> 213 def _mut_add(self, x, y):x += y; return x

UFuncTypeError: Cannot cast ufunc 'add' output from dtype('float64') to dtype('int64') with casting rule 'same_kind'
```


### System information

```shell
Name: PennyLane
Version: 0.33.1
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/PennyLaneAI/pennylane
Author:
Author-email:
License: Apache License 2.0
Location: /Users/isaac/.virtualenvs/pennylane-stable/lib/python3.9/site-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing-extensions
Required-by: PennyLane-Lightning

Platform info:           macOS-14.1.1-x86_64-i386-64bit
Python version:          3.9.14
Numpy version:           1.23.5
Scipy version:           1.11.3
Installed devices:
- default.gaussian (PennyLane-0.33.1)
- default.mixed (PennyLane-0.33.1)
- default.qubit (PennyLane-0.33.1)
- default.qubit.autograd (PennyLane-0.33.1)
- default.qubit.jax (PennyLane-0.33.1)
- default.qubit.legacy (PennyLane-0.33.1)
- default.qubit.tf (PennyLane-0.33.1)
- default.qubit.torch (PennyLane-0.33.1)
- default.qutrit (PennyLane-0.33.1)
- null.qubit (PennyLane-0.33.1)
- lightning.qubit (PennyLane-Lightning-0.33.1)
```


### Existing GitHub issues

- [X] I have searched existing GitHub issues to make sure the issue does not already exist.


Discussion:
not a solution but.. I ran this circuit with various PL versions from 0.15 up to latest and they all returned integer results 😟 

Resolving patch:
diff --git a/pennylane/measurements/sample.py b/pennylane/measurements/sample.py
--- a/pennylane/measurements/sample.py
+++ b/pennylane/measurements/sample.py
@@ -55,10 +55,37 @@ def sample(op: Optional[Union[Operator, MeasurementValue]] = None, wires=None) -
     .. note::
 
         QNodes that return samples cannot, in general, be differentiated, since the derivative
-        with respect to a sample --- a stochastic process --- is ill-defined. The one exception
-        is if the QNode uses the parameter-shift method (``diff_method="parameter-shift"``), in
-        which case ``qml.sample(obs)`` is interpreted as a single-shot expectation value of the
-        observable ``obs``.
+        with respect to a sample --- a stochastic process --- is ill-defined. An alternative
+        approach would be to use single-shot expectation values. For example, instead of this:
+
+        .. code-block:: python
+
+            dev = qml.device("default.qubit", shots=10)
+
+            @qml.qnode(dev, diff_method="parameter-shift")
+            def circuit(angle):
+                qml.RX(angle, wires=0)
+                return qml.sample(qml.PauliX(0))
+
+            angle = qml.numpy.array(0.1)
+            res = qml.jacobian(circuit)(angle)
+
+        Consider using :func:`~pennylane.expval` and a sequence of single shots, like this:
+
+        .. code-block:: python
+
+            dev = qml.device("default.qubit", shots=[(1, 10)])
+
+            @qml.qnode(dev, diff_method="parameter-shift")
+            def circuit(angle):
+                qml.RX(angle, wires=0)
+                return qml.expval(qml.PauliX(0))
+
+            def cost(angle):
+                return qml.math.hstack(circuit(angle))
+
+            angle = qml.numpy.array(0.1)
+            res = qml.jacobian(cost)(angle)
 
     **Example**
 


Test patch:
diff --git a/tests/measurements/test_sample.py b/tests/measurements/test_sample.py
--- a/tests/measurements/test_sample.py
+++ b/tests/measurements/test_sample.py
@@ -451,6 +451,36 @@ class DummyOp(Operator):  # pylint: disable=too-few-public-methods
         with pytest.raises(EigvalsUndefinedError, match="Cannot compute samples of"):
             qml.sample(op=DummyOp(0)).process_samples(samples=np.array([[1, 0]]), wire_order=[0])
 
+    def test_sample_allowed_with_parameter_shift(self):
+        """Test that qml.sample doesn't raise an error with parameter-shift and autograd."""
+        dev = qml.device("default.qubit", shots=10)
+
+        @qml.qnode(dev, diff_method="parameter-shift")
+        def circuit(angle):
+            qml.RX(angle, wires=0)
+            return qml.sample(qml.PauliX(0))
+
+        angle = qml.numpy.array(0.1)
+        res = qml.jacobian(circuit)(angle)
+        assert qml.math.shape(res) == (10,)
+        assert all(r in {-1, 0, 1} for r in np.round(res, 13))
+
+    @pytest.mark.jax
+    def test_sample_fails_with_jax_jacobian(self):
+        """Test that qml.sample raises an error with parameter-shift and jax."""
+        import jax
+
+        dev = qml.device("default.qubit", shots=10)
+
+        @qml.qnode(dev, diff_method="parameter-shift")
+        def circuit(angle):
+            qml.RX(angle, wires=0)
+            return qml.sample(qml.PauliX(0))
+
+        angle = jax.numpy.array(0.1)
+        with pytest.raises(TypeError, match=r"got int64\[10\] and float64\[10\] respectively"):
+            _ = jax.jacobian(circuit)(angle)
+
 
 @pytest.mark.jax
 @pytest.mark.parametrize("samples", (1, 10))