StatePrep accepts sparse matrices as input

[JerryChen97]

Context:

To make PennyLane use sparse matrics, at the atomic level we need to provide at least two features:

• states in PL can start from sparse
• necessary operators (e.g. QubitUnitary) have their sprase version along with the default sparse-sparse, sparse-dense multiplication

This Pull Request focus on fulfilling the first.

Note that up to v0.40, PL implements state initialization using class StatePrep from ops/qubit/state_preparation.py. Conceptually, we treated the state preparation as a generalized operation here, although unlike many other operators which are typically differentiable, StatePrep shall be expected as non-differentiable in most of the practical context.

Recall the interface of StatePrep,

class StatePrep(state, wires, pad_with=None, normalize=False, validate_norm=True)

where state is the input values for the initalization, wires is the total number of qubits of the final state, pad_with refers to a specific feature with which the user can use a shorter state of length less than 2**wires and PL will automatically fill the remaining component with the given pad_with. normalize and validate_norm are standard controls for the norm preprocessing of a state-vector.

Revisiting this, we would like to support specific input type of state as sparse matrices. Specifically, for simplicity of the first feature, we just use csr_matrix from SciPy.sparse (Compressed Sparse RowColumn format; check https://docs.scipy.org/doc/scipy/reference/sparse.html for futher details). Meanwhile, note that a nonzero pad_with is clearly against the idea of sparsity since it makes the state dense, so we will disable this option if the previous sparsity of state is given by a user.

Description of the Change:

• Introduced extra type hint and conditional branch inside the initialization of StatePrep, implementing it overloading
• Implemented the private method _sparse_statevec_permute to align with the original default StatePrep.state_vector where the input states can be embedded into the given larger wire sets with permuted wire order.
• check the parent classes for other supported operations, e.g. equal updated: out-of-scope and not necesary. Leave it to future improvement.
• improvement of the permute algorithm logic

Benefits:

• Extend the support range of StatePrep by overloading it with new features of sparse matrices.

Possible Drawbacks:

• Do not support batched states for simplicity of our MVP implementation. Open for discussion whether or not this should be included as one of the necessary requirements.
• Documentations remain missing for now, because StatePrep is actually facing public users, and it does not make sense to add notes in docs before the successful implementation of the whole workflow. They will be added in the future PR that complets the end-to-end sparse-matrix feature of PennyLane (i.e., at least a measurement can be done, which requires modification on the qubit simulator workflow)

Related GitHub Issues:
[sc-82067]

[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.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.59%. Comparing base (759dbe8) to head (397d65f).
Report is 1 commits behind head on master.

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #6863   +/-   ##
=======================================
  Coverage   99.59%   99.59%           
=======================================
  Files         480      480           
  Lines       45421    45474   +53     
=======================================
+ Hits        45236    45289   +53     
  Misses        185      185           

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

[JerryChen97]

I'll udpate changelog if we are OK with the current implementation, or until we converge at the final implementation

[AmintorDusko]

Hi @JerryChen97! Nice work! I left a few suggestions and comments.
Also, regarding your PR description, I'd like to point out that CSR stands for Compressed Sparse Row. Adding documentation from the beginning is usually a better choice, but I will let the CORE team decide.
Please feel free to tag me again whenever you're ready.
Cheers!

[AmintorDusko]

Good job!

[PietropaoloFrisoni]

🚀

[PietropaoloFrisoni]

Nothing else for me except for the last suggestions to the doc. Nice work!