Explicitly convert prob interface in `ProbsMP.process_density_matri…

…x` (#6737)

**Context:**
Previous `ProbsMP.process_density_matrix` has potential to alter the
interface. Here we convert it explicitly.

Specifically, when dealing with batched density matrices, the previous
`ProbabilityMP` did not track the interface of input `density_matrix`.

```python
import numpy as np
import pennylane as qml
from pennylane.measurements import ProbabilityMP as OriginalProbabilityMP
from pennylane.typing import TensorLike
from pennylane.wires import Wires

class CustomProbabilityMP(OriginalProbabilityMP):
    """Custom ProbabilityMP overriding process_density_matrix."""

    def process_density_matrix(self, density_matrix: TensorLike, wire_order: Wires):
        if len(np.shape(density_matrix)) == 2:
            prob = qml.math.diagonal(density_matrix)
        else:
            prob = qml.math.array(
                [qml.math.diagonal(density_matrix[i]) for i in range(np.shape(density_matrix)[0])]
            )

        # Commented out convert_like to simulate behavior without it.
        # prob = qml.math.convert_like(prob, density_matrix) 

        # Creating a pseudo-state using the diagonal probabilities
        p_state = qml.math.sqrt(prob)
        return self.process_state(p_state, wire_order)


# Instantiate the custom ProbabilityMP class
custom_prob_mp = CustomProbabilityMP()

# Define a density matrix as a NumPy array
density_matrix = qml.math.array(np.array([[[0.5, 0.5], [0.5, 0.5]], [[0.5, 0.5], [0.5, 0.5]], ]), like='torch')

# Define the wire order
wire_order = Wires([0, 1])

# Attempt to call process_density_matrix without convert_like
result = custom_prob_mp.process_density_matrix(density_matrix, wire_order)
print("Result without convert_like:", result)
# Result without convert_like: [[0.5 0.5]
# [0.5 0.5]]
```

With the fixed `ProbabilityMP` we have
```python
# Instantiate the custom ProbabilityMP class with convert_like uncommented
new_prob_mp = OriginalProbabilityMP()

# Use qml.math.convert_like internally to ensure compatibility
result = new_prob_mp.process_density_matrix(density_matrix, wire_order)
print("Result with convert_like:", result)
# Result with convert_like: tensor([[0.5000, 0.5000],
#        [0.5000, 0.5000]], dtype=torch.float64)
```

**Description of the Change:**
Enforce the probs to match input interface via `qml.math.convert_like`

**Benefits:**
This PR is blocking #6684.
Its merge into master will free new `default.mixed` from many roundabout
solutions.

**Possible Drawbacks:**

**Related GitHub Issues:**

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}"