Explicitly convert prob interface in ProbsMP.process_density_matrix

pennylane/measurements/probs.py

@@ -264,6 +264,7 @@ def process_density_matrix(self, density_matrix: TensorLike, wire_order: Wires):
             )
 
         # Since we only care about the probabilities, we can simplify the task here by creating a 'pseudo-state' to carry the diagonal elements and reuse the process_state method
+        prob = qml.math.convert_like(prob, density_matrix)
         p_state = qml.math.sqrt(prob)
         return self.process_state(p_state, wire_order)
 

tests/measurements/test_probs.py

@@ -205,6 +205,7 @@ def test_process_density_matrix_basic(self, interface):
         wires = qml.wires.Wires(range(1))
         expected = qml.math.array([0.5, 0.5], like=interface)
         calculated_probs = qml.probs().process_density_matrix(dm, wires)
+        assert qml.math.get_interface(calculated_probs) == interface
         assert qml.math.allclose(calculated_probs, expected)
 
     @pytest.mark.all_interfaces
@@ -226,6 +227,7 @@ def test_process_density_matrix_subsets(self, interface, subset_wires, expected)
         )
         wires = qml.wires.Wires(range(2))
         subset_probs = qml.probs(wires=subset_wires).process_density_matrix(dm, wires)
+        assert qml.math.get_interface(subset_probs) == interface
         assert subset_probs.shape == qml.math.shape(expected)
         assert qml.math.allclose(subset_probs, expected)
 
@@ -261,6 +263,7 @@ def test_process_density_matrix_batched(self, interface, subset_wires, expected)
 
         expected = qml.math.array(expected, like=interface)
         # Check if the calculated probabilities match the expected values
+        assert qml.math.get_interface(subset_probs) == interface
         assert (
             subset_probs.shape == expected.shape
         ), f"Shape mismatch: expected {expected.shape}, got {subset_probs.shape}"
@@ -307,6 +310,7 @@ def test_process_density_matrix_medium(self, interface, subset_wires):
         expected = np.diag(reduced_dm)
         expected = qml.math.array(expected, like=interface)
         # Check if the calculated probabilities match the expected values
+        assert qml.math.get_interface(subset_probs) == interface
         assert (
             subset_probs.shape == expected.shape
         ), f"Shape mismatch: expected {expected.shape}, got {subset_probs.shape}"