diff --git a/pennylane/measurements/probs.py b/pennylane/measurements/probs.py index 77c9f61e4df..5e5126e5fb0 100644 --- a/pennylane/measurements/probs.py +++ b/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) diff --git a/tests/measurements/test_probs.py b/tests/measurements/test_probs.py index 6254c1be12b..07bd8fcd463 100644 --- a/tests/measurements/test_probs.py +++ b/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}"