Explicitly convert prob interface in ProbsMP.process_density_matrix

[JerryChen97]

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.

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

# 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:

[github-actions]

Hello. You may have forgotten to update the changelog!
Please edit doc/releases/changelog-dev.md with:

• A one-to-two sentence description of the change. You may include a small working example for new features.
• A link back to this PR.
• Your name (or GitHub username) in the contributors section.

[PietropaoloFrisoni]

Thanks @JerryChen97.

I just have one question: why create a single 1-line PR for this instead of applying the change directly in #6684 ?

[andrijapau]

We should probably add a test to ensure desired behaviour (interface remaining the same)?

[JerryChen97]

Thanks @JerryChen97.

I just have one question: why create a single 1-line PR for this instead of applying the change directly in #6684 ?

The modification to ProbsMP is of different scope from #6684.
I made same modification there but I dont expect it to merge within 2weeks.
However, ProbsMP as its own can be used in any ways allowed now, which contain this same bug and we should fix it sasp

[PietropaoloFrisoni]

@JerryChen97 Thanks. Then I would say that a test and possibly a changelog entry would be ideal here 👍

Can you also update the PR description with an example of a non-working workflow that is solved by this bug fix?

[JerryChen97]

@JerryChen97 Thanks. Then I would say that a test and possibly a changelog entry would be ideal here 👍

Can you also update the PR description with an example of a non-working workflow that is solved by this bug fix?

Done

[JerryChen97]

result = custom_prob_mp.process_density_matrix(density_matrix, wire_order)

Added interface assert for all the current 4 different testing branches

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.60%. Comparing base (585767c) to head (922f808).
Report is 3 commits behind head on master.

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #6737   +/-   ##
=======================================
  Coverage   99.60%   99.60%           
=======================================
  Files         475      475           
  Lines       45102    45118   +16     
=======================================
+ Hits        44925    44941   +16     
  Misses        177      177           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[andrijapau]

Nice, thanks for adding the tests 🎉