Remove op.is_hermitian check in expval, counts, sample to allow jit tracing

doc/releases/changelog-dev.md

@@ -258,6 +258,9 @@
   (which is not currently compatible with `KerasLayer`), linking to instructions to enable Keras 2.
   [(#5488)](https://github.com/PennyLaneAI/pennylane/pull/5488)
 
+* Removed the warning that an observable might not be hermitian in `qnode` executions. This enables jit-compilation.
+  [(#5506)](https://github.com/PennyLaneAI/pennylane/pull/5506)
+
 <h3>Breaking changes 💔</h3>
 
 * The private functions `_pauli_mult`, `_binary_matrix` and `_get_pauli_map` from the `pauli` module have been removed. The same functionality can be achieved using newer features in the ``pauli`` module.

pennylane/measurements/counts.py

@@ -14,7 +14,6 @@
 """
 This module contains the qml.counts measurement.
 """
-import warnings
 from typing import Sequence, Tuple, Optional
 import numpy as np
 
@@ -146,9 +145,6 @@ def circuit():
 
         return CountsMP(obs=op, all_outcomes=all_outcomes)
 
-    if op is not None and not op.is_hermitian:  # None type is also allowed for op
-        warnings.warn(f"{op.name} might not be hermitian.")
-
     if wires is not None:
         if op is not None:
             raise ValueError(

pennylane/measurements/expval.py

@@ -14,7 +14,6 @@
 """
 This module contains the qml.expval measurement.
 """
-import warnings
 from typing import Sequence, Tuple, Union
 
 import pennylane as qml
@@ -71,9 +70,6 @@ def circuit(x):
             "Expectation values of qml.Identity() without wires are currently not allowed."
         )
 
-    if not op.is_hermitian:
-        warnings.warn(f"{op.name} might not be hermitian.")
-
     return ExpectationMP(obs=op)
 
 

pennylane/measurements/sample.py

@@ -15,7 +15,6 @@
 This module contains the qml.sample measurement.
 """
 import functools
-import warnings
 from typing import Sequence, Tuple, Optional, Union
 
 import numpy as np
@@ -173,9 +172,6 @@ def __init__(self, obs=None, wires=None, eigvals=None, id=None):
             super().__init__(obs=obs)
             return
 
-        if obs is not None and not obs.is_hermitian:  # None type is also allowed for op
-            warnings.warn(f"{obs.name} might not be hermitian.")
-
         if wires is not None:
             if obs is not None:
                 raise ValueError(

pennylane/measurements/var.py

@@ -15,7 +15,6 @@
 """
 This module contains the qml.var measurement.
 """
-import warnings
 from typing import Sequence, Tuple, Union
 
 import pennylane as qml
@@ -64,9 +63,6 @@ def circuit(x):
             "qml.var does not support measuring sequences of measurements or observables"
         )
 
-    if not op.is_hermitian:
-        warnings.warn(f"{op.name} might not be hermitian.")
-
     return VarianceMP(obs=op)
 
 

tests/devices/default_qubit/test_default_qubit.py

@@ -2052,6 +2052,34 @@ def circuit():
 
         assert circuit() == expected
 
+    @pytest.mark.jax
+    @pytest.mark.parametrize("measurement_func", [qml.expval, qml.var])
+    def test_differentiate_jitted_qnode(self, measurement_func):
+        """Test that a jitted qnode can be correctly differentiated"""
+        import jax
+
+        dev = DefaultQubit()
+
+        def qfunc(x, y):
+            qml.RX(x, 0)
+            return measurement_func(qml.Hamiltonian(y, [qml.Z(0)]))
+
+        qnode = qml.QNode(qfunc, dev, interface="jax")
+        qnode_jit = jax.jit(qml.QNode(qfunc, dev, interface="jax"))
+
+        x = jax.numpy.array(0.5)
+        y = jax.numpy.array([0.5])
+
+        res = qnode(x, y)
+        res_jit = qnode_jit(x, y)
+
+        assert qml.math.allclose(res, res_jit)
+
+        grad = jax.grad(qnode)(x, y)
+        grad_jit = jax.grad(qnode_jit)(x, y)
+
+        assert qml.math.allclose(grad, grad_jit)
+
 
 @pytest.mark.parametrize("max_workers", max_workers_list)
 def test_broadcasted_parameter(max_workers):

tests/interfaces/default_qubit_2_integration/test_jax_jit_qnode_default_qubit_2.py

@@ -1520,6 +1520,7 @@ def test_hamiltonian_expansion_analytic(
         spy = mocker.spy(qml.transforms, "hamiltonian_expand")
         obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)]
 
+        @jax.jit
         @qnode(
             dev,
             interface=interface,

tests/measurements/legacy/test_expval_legacy.py

@@ -80,24 +80,6 @@ def circuit(x):
 
         custom_measurement_process(new_dev, spy)
 
-    def test_not_an_observable(self, mocker):
-        """Test that a warning is raised if the provided
-        argument might not be hermitian."""
-        dev = qml.device("default.qubit.legacy", wires=2)
-
-        @qml.qnode(dev)
-        def circuit():
-            qml.RX(0.52, wires=0)
-            return qml.expval(qml.prod(qml.PauliX(0), qml.PauliZ(0)))
-
-        new_dev = circuit.device
-        spy = mocker.spy(qml.QubitDevice, "expval")
-
-        with pytest.warns(UserWarning, match="Prod might not be hermitian."):
-            _ = circuit()
-
-        custom_measurement_process(new_dev, spy)
-
     def test_observable_return_type_is_expectation(self, mocker):
         """Test that the return type of the observable is :attr:`ObservableReturnTypes.Expectation`"""
         dev = qml.device("default.qubit.legacy", wires=2)

tests/measurements/legacy/test_measurements_legacy.py

@@ -24,8 +24,6 @@
     Shots,
     StateMeasurement,
     StateMP,
-    expval,
-    var,
 )
 
 from pennylane.wires import Wires
@@ -66,22 +64,6 @@ def test_shape_unrecognized_error():
         mp.shape(dev, Shots(None))
 
 
-@pytest.mark.parametrize("stat_func", [expval, var])
-def test_not_an_observable(stat_func):
-    """Test that a UserWarning is raised if the provided
-    argument might not be hermitian."""
-
-    dev = qml.device("default.qubit.legacy", wires=2)
-
-    @qml.qnode(dev)
-    def circuit():
-        qml.RX(0.52, wires=0)
-        return stat_func(qml.prod(qml.PauliX(0), qml.PauliZ(0)))
-
-    with pytest.warns(UserWarning, match="Prod might not be hermitian."):
-        _ = circuit()
-
-
 class TestSampleMeasurement:
     """Tests for the SampleMeasurement class."""
 

tests/measurements/legacy/test_sample_legacy.py

@@ -251,22 +251,6 @@ def circuit():
 
         custom_measurement_process(dev, spy)
 
-    def test_not_an_observable(self, mocker):
-        """Test that a UserWarning is raised if the provided
-        argument might not be hermitian."""
-        dev = qml.device("default.qubit.legacy", wires=2, shots=10)
-        spy = mocker.spy(qml.QubitDevice, "sample")
-
-        @qml.qnode(dev)
-        def circuit():
-            qml.RX(0.52, wires=0)
-            return qml.sample(qml.prod(qml.PauliX(0), qml.PauliZ(0)))
-
-        with pytest.warns(UserWarning, match="Prod might not be hermitian."):
-            _ = circuit()
-
-        custom_measurement_process(dev, spy)
-
     def test_observable_return_type_is_sample(self, mocker):
         """Test that the return type of the observable is :attr:`ObservableReturnTypes.Sample`"""
         n_shots = 10

tests/measurements/legacy/test_var_legacy.py

@@ -77,22 +77,6 @@ def circuit(x):
 
         custom_measurement_process(dev, spy)
 
-    def test_not_an_observable(self, mocker):
-        """Test that a UserWarning is raised if the provided
-        argument might not be hermitian."""
-        dev = qml.device("default.qubit.legacy", wires=2)
-        spy = mocker.spy(qml.QubitDevice, "var")
-
-        @qml.qnode(dev)
-        def circuit():
-            qml.RX(0.52, wires=0)
-            return qml.var(qml.prod(qml.PauliX(0), qml.PauliZ(0)))
-
-        with pytest.warns(UserWarning, match="Prod might not be hermitian."):
-            _ = circuit()
-
-        custom_measurement_process(dev, spy)
-
     def test_observable_return_type_is_variance(self, mocker):
         """Test that the return type of the observable is :attr:`ObservableReturnTypes.Variance`"""
         dev = qml.device("default.qubit.legacy", wires=2)

tests/measurements/test_counts.py

@@ -61,13 +61,6 @@ def test_providing_observable_and_wires(self):
         ):
             qml.counts(qml.PauliZ(0), wires=[0, 1])
 
-    def test_observable_might_not_be_hermitian(self):
-        """Test that a UserWarning is raised if the provided
-        argument might not be hermitian."""
-
-        with pytest.warns(UserWarning, match="Prod might not be hermitian."):
-            qml.counts(qml.prod(qml.PauliX(0), qml.PauliZ(0)))
-
     def test_hash(self):
         """Test that the hash property includes the all_outcomes property."""
         m1 = qml.counts(all_outcomes=True)

tests/measurements/test_expval.py

@@ -73,19 +73,6 @@ def circuit(x):
         else:
             assert res.dtype == r_dtype
 
-    def test_not_an_observable(self):
-        """Test that a warning is raised if the provided
-        argument might not be hermitian."""
-        dev = qml.device("default.qubit", wires=2)
-
-        @qml.qnode(dev)
-        def circuit():
-            qml.RX(0.52, wires=0)
-            return qml.expval(qml.prod(qml.PauliX(0), qml.PauliZ(0)))
-
-        with pytest.warns(UserWarning, match="Prod might not be hermitian."):
-            _ = circuit()
-
     def test_observable_return_type_is_expectation(self):
         """Test that the return type of the observable is :attr:`ObservableReturnTypes.Expectation`"""
         dev = qml.device("default.qubit", wires=2)

tests/measurements/test_measurements.py

@@ -293,22 +293,6 @@ def test_queueing_tensor_observable(self, op1, op2, stat_func, return_type):
         assert isinstance(meas_proc, MeasurementProcess)
         assert meas_proc.return_type == return_type
 
-    def test_not_an_observable(self, stat_func, return_type):  # pylint: disable=unused-argument
-        """Test that a UserWarning is raised if the provided
-        argument might not be hermitian."""
-        if stat_func is sample:
-            pytest.skip("Sampling is not yet supported with symbolic operators.")
-
-        dev = qml.device("default.qubit", wires=2)
-
-        @qml.qnode(dev)
-        def circuit():
-            qml.RX(0.52, wires=0)
-            return stat_func(qml.prod(qml.PauliX(0), qml.PauliZ(0)))
-
-        with pytest.warns(UserWarning, match="Prod might not be hermitian."):
-            _ = circuit()
-
 
 class TestProperties:
     """Test for the properties"""

tests/measurements/test_sample.py

@@ -126,19 +126,6 @@ def circuit():
         assert result[2].dtype == np.dtype("int")
         assert np.array_equal(result[2].shape, (n_sample,))
 
-    def test_not_an_observable(self):
-        """Test that a UserWarning is raised if the provided
-        argument might not be hermitian."""
-        dev = qml.device("default.qubit", wires=2, shots=10)
-
-        @qml.qnode(dev)
-        def circuit():
-            qml.RX(0.52, wires=0)
-            return qml.sample(qml.prod(qml.PauliX(0), qml.PauliZ(0)))
-
-        with pytest.warns(UserWarning, match="Prod might not be hermitian."):
-            _ = circuit()
-
     def test_observable_return_type_is_sample(self):
         """Test that the return type of the observable is :attr:`ObservableReturnTypes.Sample`"""
         n_shots = 10

tests/measurements/test_var.py

@@ -49,19 +49,6 @@ def circuit(x):
         else:
             assert res.dtype == r_dtype
 
-    def test_not_an_observable(self):
-        """Test that a UserWarning is raised if the provided
-        argument might not be hermitian."""
-        dev = qml.device("default.qubit", wires=2)
-
-        @qml.qnode(dev)
-        def circuit():
-            qml.RX(0.52, wires=0)
-            return qml.var(qml.prod(qml.PauliX(0), qml.PauliZ(0)))
-
-        with pytest.warns(UserWarning, match="Prod might not be hermitian."):
-            _ = circuit()
-
     def test_observable_return_type_is_variance(self):
         """Test that the return type of the observable is :attr:`ObservableReturnTypes.Variance`"""
         dev = qml.device("default.qubit", wires=2)

tests/ops/op_math/test_prod.py

@@ -1485,8 +1485,8 @@ def circuit(weights):
         true_grad = -qnp.sqrt(2) * qnp.cos(weights[0] / 2) * qnp.sin(weights[0] / 2)
         assert qnp.allclose(grad, true_grad)
 
-    def test_non_hermitian_obs_not_supported(self):
-        """Test that non-hermitian ops in a measurement process will raise a warning."""
+    def test_non_supported_obs_not_supported(self):
+        """Test that non-supported ops in a measurement process will raise an error."""
         wires = [0, 1]
         dev = qml.device("default.qubit", wires=wires)
         prod_op = Prod(qml.RX(1.23, wires=0), qml.Identity(wires=1))

tests/ops/op_math/test_sprod.py

@@ -1084,20 +1084,6 @@ def circuit(weights):
         true_grad = 100 * -qnp.sqrt(2) * qnp.cos(weights[0] / 2) * qnp.sin(weights[0] / 2)
         assert qnp.allclose(grad, true_grad)
 
-    def test_non_hermitian_obs_not_supported(self):
-        """Test that non-hermitian ops in a measurement process will raise a warning."""
-        wires = [0, 1]
-        dev = qml.device("default.qubit", wires=wires)
-        sprod_op = SProd(1.0 + 2.0j, qml.RX(1.23, wires=0))
-
-        @qml.qnode(dev)
-        def my_circ():
-            qml.PauliX(0)
-            return qml.expval(sprod_op)
-
-        with pytest.raises(NotImplementedError):
-            my_circ()
-
     @pytest.mark.torch
     @pytest.mark.parametrize("diff_method", ("parameter-shift", "backprop"))
     def test_torch(self, diff_method):

tests/ops/op_math/test_sum.py

@@ -1267,20 +1267,6 @@ def circuit(weights):
         true_grad = qnp.array([-0.09347337, -0.18884787, -0.28818254])
         assert qnp.allclose(grad, true_grad)
 
-    def test_non_hermitian_op_in_measurement_process(self):
-        """Test that non-hermitian ops in a measurement process will raise a warning."""
-        wires = [0, 1]
-        dev = qml.device("default.qubit", wires=wires)
-        sum_op = Sum(Prod(qml.RX(1.23, wires=0), qml.Identity(wires=1)), qml.Identity(wires=1))
-
-        @qml.qnode(dev, interface=None)
-        def my_circ():
-            qml.PauliX(0)
-            return qml.expval(sum_op)
-
-        with pytest.warns(UserWarning, match="Sum might not be hermitian."):
-            my_circ()
-
     def test_params_can_be_considered_trainable(self):
         """Tests that the parameters of a Sum are considered trainable."""
         dev = qml.device("default.qubit", wires=2)