qml.state and qml.density_matrix custom wires (#2779)

* remove skips now that default.mixed supports backprop

* remove custom wire labels error; tests

* changelog

* more test cases

doc/releases/changelog-dev.md

@@ -48,6 +48,10 @@
 
 <h3>Improvements</h3>
 
+* The `qml.state` and `qml.density_matrix` measurements now support custom wire
+  labels.
+  [(#2779)](https://github.com/PennyLaneAI/pennylane/pull/2779)
+
 * Adds a new function to compare operators. `qml.equal` can be used to compare equality of parametric operators taking into account their interfaces and trainability.
   [(#2651)](https://github.com/PennyLaneAI/pennylane/pull/2651)
 

pennylane/_qubit_device.py

@@ -488,10 +488,7 @@ def statistics(self, observables, shot_range=None, bin_size=None):
                         "The state or density matrix cannot be returned in combination"
                         " with other return types"
                     )
-                if self.wires.labels != tuple(range(self.num_wires)):
-                    raise qml.QuantumFunctionError(
-                        "Returning the state is not supported when using custom wire labels"
-                    )
+
                 # Check if the state is accessible and decide to return the state or the density
                 # matrix.
                 results.append(self.access_state(wires=obs.wires))
@@ -681,6 +678,7 @@ def density_matrix(self, wires):
             representing the reduced density matrix of the state prior to measurement.
         """
         state = getattr(self, "state", None)
+        wires = self.map_wires(wires)
         return qml.math.reduced_dm(state, indices=wires, c_dtype=self.C_DTYPE)
 
     def vn_entropy(self, wires, log_base):

tests/devices/test_default_mixed_autograd.py

@@ -315,13 +315,14 @@ def circuit(x):
         res = grad_fn(p)
         assert np.allclose(res, qml.jacobian(circuit2)(p), atol=tol, rtol=0)
 
-    def test_state_differentiability(self, tol):
+    @pytest.mark.parametrize("wires", [[0], ["abc"]])
+    def test_state_differentiability(self, wires, tol):
         """Test that the device state can be differentiated"""
-        dev = qml.device("default.mixed", wires=1)
+        dev = qml.device("default.mixed", wires=wires)
 
         @qml.qnode(dev, diff_method="backprop", interface="autograd")
         def circuit(a):
-            qml.RY(a, wires=0)
+            qml.RY(a, wires=wires[0])
             return qml.state()
 
         a = np.array(0.54, requires_grad=True)
@@ -337,15 +338,16 @@ def cost(a):
         expected = np.sin(a)
         assert np.allclose(grad, expected, atol=tol, rtol=0)
 
-    def test_density_matrix_differentiability(self, tol):
+    @pytest.mark.parametrize("wires", [range(2), [-12.32, "abc"]])
+    def test_density_matrix_differentiability(self, wires, tol):
         """Test that the density matrix can be differentiated"""
-        dev = qml.device("default.mixed", wires=2)
+        dev = qml.device("default.mixed", wires=wires)
 
         @qml.qnode(dev, diff_method="backprop", interface="autograd")
         def circuit(a):
-            qml.RY(a, wires=0)
-            qml.CNOT(wires=[0, 1])
-            return qml.density_matrix(wires=1)
+            qml.RY(a, wires=wires[0])
+            qml.CNOT(wires=[wires[0], wires[1]])
+            return qml.density_matrix(wires=wires[1])
 
         a = np.array(0.54, requires_grad=True)
 

tests/devices/test_default_mixed_jax.py

@@ -329,16 +329,17 @@ def circuit(a):
 
         assert np.allclose(grad, expected, atol=tol, rtol=0)
 
+    @pytest.mark.parametrize("wires", [range(2), [-12.32, "abc"]])
     @pytest.mark.parametrize("decorator", decorators)
-    def test_density_matrix_differentiability(self, decorator, tol):
+    def test_density_matrix_differentiability(self, decorator, wires, tol):
         """Test that the density matrix can be differentiated"""
-        dev = qml.device("default.mixed", wires=2)
+        dev = qml.device("default.mixed", wires=wires)
 
         @qml.qnode(dev, diff_method="backprop", interface="jax")
         def circuit(a):
-            qml.RY(a, wires=0)
-            qml.CNOT(wires=[0, 1])
-            return qml.density_matrix(wires=1)
+            qml.RY(a, wires=wires[0])
+            qml.CNOT(wires=[wires[0], wires[1]])
+            return qml.density_matrix(wires=wires[1])
 
         a = jnp.array(0.54)
 

tests/devices/test_default_mixed_tf.py

@@ -298,15 +298,16 @@ def circuit(x):
         res = tape.jacobian(res, p_tf)
         assert np.allclose(res, qml.jacobian(circuit2)(p), atol=tol, rtol=0)
 
+    @pytest.mark.parametrize("wires", [[0], ["abc"]])
     @pytest.mark.parametrize("decorator, interface", decorators_interfaces)
-    def test_state_differentiability(self, decorator, interface, tol):
+    def test_state_differentiability(self, decorator, interface, wires, tol):
         """Test that the device state can be differentiated"""
-        dev = qml.device("default.mixed", wires=1)
+        dev = qml.device("default.mixed", wires=wires)
 
         @decorator
         @qml.qnode(dev, interface=interface, diff_method="backprop")
         def circuit(a):
-            qml.RY(a, wires=0)
+            qml.RY(a, wires=wires[0])
             return qml.state()
 
         a = tf.Variable(0.54, trainable=True)
@@ -343,16 +344,17 @@ def circuit(a):
         assert np.allclose(grad, expected, atol=tol, rtol=0)
 
     @pytest.mark.parametrize("decorator, interface", decorators_interfaces)
-    def test_density_matrix_differentiability(self, decorator, interface, tol):
+    @pytest.mark.parametrize("wires", [range(2), [-12.32, "abc"]])
+    def test_density_matrix_differentiability(self, decorator, interface, wires, tol):
         """Test that the density matrix can be differentiated"""
-        dev = qml.device("default.mixed", wires=2)
+        dev = qml.device("default.mixed", wires=wires)
 
         @decorator
         @qml.qnode(dev, diff_method="backprop", interface=interface)
         def circuit(a):
-            qml.RY(a, wires=0)
-            qml.CNOT(wires=[0, 1])
-            return qml.density_matrix(wires=1)
+            qml.RY(a, wires=wires[0])
+            qml.CNOT(wires=[wires[0], wires[1]])
+            return qml.density_matrix(wires=wires[1])
 
         a = tf.Variable(0.54, trainable=True)
 

tests/devices/test_default_mixed_torch.py

@@ -286,13 +286,14 @@ def circuit(x):
         grad = torch.autograd.functional.jacobian(circuit1, p_torch)
         assert qml.math.allclose(grad, qml.jacobian(circuit2)(p), atol=tol, rtol=0)
 
-    def test_state_differentiability(self, tol):
+    @pytest.mark.parametrize("wires", [[0], ["abc"]])
+    def test_state_differentiability(self, wires, tol):
         """Test that the device state can be differentiated"""
-        dev = qml.device("default.mixed", wires=1)
+        dev = qml.device("default.mixed", wires=wires)
 
         @qml.qnode(dev, diff_method="backprop", interface="torch")
         def circuit(a):
-            qml.RY(a, wires=0)
+            qml.RY(a, wires=wires[0])
             return qml.state()
 
         a = torch.tensor(0.54, dtype=torch.float64, requires_grad=True)
@@ -321,15 +322,16 @@ def circuit(a):
 
         assert np.allclose(grad, expected, atol=tol, rtol=0)
 
-    def test_density_matrix_differentiability(self, tol):
+    @pytest.mark.parametrize("wires", [range(2), [-12.32, "abc"]])
+    def test_density_matrix_differentiability(self, wires, tol):
         """Test that the density matrix can be differentiated"""
-        dev = qml.device("default.mixed", wires=2)
+        dev = qml.device("default.mixed", wires=wires)
 
         @qml.qnode(dev, diff_method="backprop", interface="torch")
         def circuit(a):
-            qml.RY(a, wires=0)
-            qml.CNOT(wires=[0, 1])
-            return qml.density_matrix(wires=1)
+            qml.RY(a, wires=wires[0])
+            qml.CNOT(wires=[wires[0], wires[1]])
+            return qml.density_matrix(wires=wires[1])
 
         a = torch.tensor(0.54, dtype=torch.float64, requires_grad=True)
 

tests/devices/test_default_qubit_autograd.py

@@ -264,22 +264,22 @@ def circuit(x):
         res = grad_fn(p)
         assert np.allclose(res, qml.jacobian(circuit2)(p), atol=tol, rtol=0)
 
-    def test_state_differentiability(self, tol):
+    @pytest.mark.parametrize("wires", [[0], ["abc"]])
+    def test_state_differentiability(self, wires, tol):
         """Test that the device state can be differentiated"""
-        dev = qml.device("default.qubit.autograd", wires=1)
+        dev = qml.device("default.qubit.autograd", wires=wires)
 
         @qml.qnode(dev, diff_method="backprop", interface="autograd")
         def circuit(a):
-            qml.RY(a, wires=0)
-            return qml.expval(qml.PauliZ(0))
+            qml.RY(a, wires=wires[0])
+            return qml.state()
 
         a = np.array(0.54, requires_grad=True)
 
         def cost(a):
             """A function of the device quantum state, as a function
             of input QNode parameters."""
-            circuit(a)
-            res = np.abs(dev.state) ** 2
+            res = np.abs(circuit(a)) ** 2
             return res[1] - res[0]
 
         grad = qml.grad(cost)(a)

tests/devices/test_default_qubit_jax.py

@@ -446,13 +446,14 @@ def circuit(x):
         )
         assert jnp.allclose(res, expected, atol=tol, rtol=0)
 
-    def test_state_differentiability(self, tol):
+    @pytest.mark.parametrize("wires", [[0], ["abc"]])
+    def test_state_differentiability(self, wires, tol):
         """Test that the device state can be differentiated"""
-        dev = qml.device("default.qubit.jax", wires=1)
+        dev = qml.device("default.qubit.jax", wires=wires)
 
         @qml.qnode(dev, diff_method="backprop", interface="jax")
         def circuit(a):
-            qml.RY(a, wires=0)
+            qml.RY(a, wires=wires[0])
             return qml.state()
 
         a = jnp.array(0.54)

tests/devices/test_default_qubit_tf.py

@@ -1268,20 +1268,20 @@ def circuit(x):
         res = tape.jacobian(res, p_tf)
         assert np.allclose(res, qml.jacobian(circuit2)(p), atol=tol, rtol=0)
 
-    def test_state_differentiability(self, tol):
+    @pytest.mark.parametrize("wires", [[0], ["abc"]])
+    def test_state_differentiability(self, wires, tol):
         """Test that the device state can be differentiated"""
-        dev = qml.device("default.qubit.tf", wires=1)
+        dev = qml.device("default.qubit.tf", wires=wires)
 
         @qml.qnode(dev, diff_method="backprop", interface="tf")
         def circuit(a):
-            qml.RY(a, wires=0)
-            return qml.expval(qml.PauliZ(0))
+            qml.RY(a, wires=wires[0])
+            return qml.state()
 
         a = tf.Variable(0.54)
 
         with tf.GradientTape() as tape:
-            circuit(a)
-            res = tf.abs(dev.state) ** 2
+            res = tf.abs(circuit(a)) ** 2
             res = res[1] - res[0]
 
         grad = tape.gradient(res, a)

tests/devices/test_default_qubit_torch.py

@@ -1408,19 +1408,19 @@ def circuit(x):
         p_grad = p_torch.grad
         assert qml.math.allclose(p_grad, qml.jacobian(circuit2)(p), atol=tol, rtol=0)
 
-    def test_state_differentiability(self, device, torch_device, tol):
+    @pytest.mark.parametrize("wires", [[0], ["abc"]])
+    def test_state_differentiability(self, device, torch_device, wires, tol):
         """Test that the device state can be differentiated"""
-        dev = qml.device("default.qubit.torch", wires=1, torch_device=torch_device)
+        dev = qml.device("default.qubit.torch", wires=wires, torch_device=torch_device)
 
         @qml.qnode(dev, diff_method="backprop", interface="torch")
         def circuit(a):
-            qml.RY(a, wires=0)
-            return qml.expval(qml.PauliZ(0))
+            qml.RY(a, wires=wires[0])
+            return qml.state()
 
         a = torch.tensor(0.54, requires_grad=True, device=torch_device)
 
-        circuit(a)
-        res = torch.abs(dev.state) ** 2
+        res = torch.abs(circuit(a)) ** 2
         res = res[1] - res[0]
         res.backward()
 

tests/test_measurements.py

@@ -1143,18 +1143,19 @@ def loss_fn(x):
 
     @pytest.mark.parametrize("wires", [[0, 2, 3, 1], ["a", -1, "b", 1000]])
     def test_custom_wire_labels(self, wires):
-        """Test if an error is raised when custom wire labels are used"""
+        """Test the state when custom wire labels are used"""
         dev = qml.device("default.qubit", wires=wires)
 
         @qml.qnode(dev, diff_method="parameter-shift")
         def func():
-            qml.Hadamard(wires=wires[0])
-            for i in range(3):
-                qml.CNOT(wires=[wires[i], wires[i + 1]])
+            for i in range(4):
+                qml.Hadamard(wires[i])
             return state()
 
-        with pytest.raises(qml.QuantumFunctionError, match="custom wire labels"):
-            func()
+        state_expected = 0.25 * np.ones(16)
+        state_val = func()
+
+        assert np.allclose(state_expected, state_val)
 
     @pytest.mark.parametrize("shots", [None, 1, 10])
     def test_shape(self, shots):
@@ -1450,21 +1451,50 @@ def func():
         with pytest.raises(qml.QuantumFunctionError, match="Returning the state is not supported"):
             func()
 
-    @pytest.mark.parametrize("wires", [[0, 2, 3, 1], ["a", -1, "b", 1000]])
+    @pytest.mark.parametrize("wires", [[0, 2], ["a", -1]])
     @pytest.mark.parametrize("dev_name", ["default.qubit", "default.mixed"])
     def test_custom_wire_labels(self, wires, dev_name):
-        """Test if an error is raised when custom wire labels are used"""
+        """Test that the correct density matrix for an example with a mixed
+        state when using custom wires"""
+
         dev = qml.device(dev_name, wires=wires)
 
-        @qml.qnode(dev, diff_method="parameter-shift")
+        @qml.qnode(dev)
         def func():
-            qml.Hadamard(wires=wires[0])
-            for i in range(3):
-                qml.CNOT(wires=[wires[i], wires[i + 1]])
-            return density_matrix(0)
+            qml.Hadamard(wires[0])
+            qml.CNOT(wires=[wires[0], wires[1]])
+            return qml.density_matrix(wires=wires[1])
 
-        with pytest.raises(qml.QuantumFunctionError, match="custom wire labels"):
-            func()
+        density = func()
+
+        assert np.allclose(np.array([[0.5 + 0.0j, 0.0 + 0.0j], [0.0 + 0.0j, 0.5 + 0.0j]]), density)
+
+    @pytest.mark.parametrize("wires", [[3, 1], ["b", 1000]])
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.mixed"])
+    def test_custom_wire_labels_all_wires(self, wires, dev_name):
+        """Test that the correct density matrix for an example with a mixed
+        state when using custom wires"""
+        dev = qml.device(dev_name, wires=wires)
+
+        @qml.qnode(dev)
+        def func():
+            qml.Hadamard(wires[0])
+            qml.CNOT(wires=[wires[0], wires[1]])
+            return qml.density_matrix(wires=[wires[0], wires[1]])
+
+        density = func()
+
+        assert np.allclose(
+            np.array(
+                [
+                    [0.5 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.5 + 0.0j],
+                    [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
+                    [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
+                    [0.5 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.5 + 0.0j],
+                ]
+            ),
+            density,
+        )
 
     @pytest.mark.parametrize("shots", [None, 1, 10])
     def test_shape(self, shots):

tests/transforms/test_condition.py

@@ -37,6 +37,7 @@
     qml.var(qml.PauliX("b")),
     qml.state(),
     qml.density_matrix(wires=[2, 3]),
+    qml.density_matrix(wires=["b", -231]),
 ]