Repository: PennyLaneAI/pennylane
Pull request #6889 (2025-01-27 19:49:27)

`QubitUnitary` accepts sparse matrices as inputs
**Context:**
As a follow-up of https://github.com/PennyLaneAI/pennylane/pull/6863, in this PR we would like to introduce the overloaded `QubitUnitary` capable of taking sparse matrices and processing along.

Currently, if a user create a new `QubitUnitary` with sparse matrix as input, it will succeed with `unitary_check=False` but fail with `True`. This is because lots of the internal processing make use of incompatible functionalities e.g. `einsum` for `csr_matrix`, for which we will create more sparse-specific routines to help build a reasonable workflow for the sparse cases.

**Description of the Change:**
`QubitUnitary` has quite some API that we need to consider, some of which might need further attention to decide whether or not to implement for sparse matrices right now.
 - [x] `init`: inserted extra conditional unitary check, and separated the previous logic as static `_unitary_check`
 - [x]  `matrix()`: the method to generate `QubitUnitary`'s matrix should be able to return a sparse as well
 - [x] `adjoint`: conditional treatment for sparse matrices, since our default dispatch method does not work for `csr_matrix`
 - [x] `pow`: similar as above, although this could be easily merged into our multi-dispatch in the future
 - [x] `compute_decomposition`: to avoid unnecessary extra work, one- and two-qubit sparse ops are still converted back to desnse matrices, although in practice it actually does not even make sense to use sparse matrices at this level (recalling our need for sparse matrices, it's for systems of size ~20)
 - [x] support for wires of size 20: should be able to init a sparse `QubitUnitary` obj with 20 wires.

**Benefits:**
Necessary extension for `QubitUnitary` implemented.

**Possible Drawbacks:**
Implementations are quite scattered in the source code. In the future we can think about adding overloading to `qml.math` instead to help unify the matrix manipulations of sparse matrices; this is needed for `QubitUnitary` to act like a 'normal' operator as others, or it won't have enough capabilities, e.g. `QU @ QU` won't work if both implemented as `csr_matrix`, because PL is confulsed about how to dispatch for CSR's when it comes to necessary math methods e.g. `reshape`. These are going to be added in [another PR](https://github.com/PennyLaneAI/pennylane/pull/6947)

Also, to minimize the amount of conditional branches specifically for `csr_matrix`, we neglected the following functionalities:
 - `_controlled`: this method actually involves the sparse implementation of `ControlledQubitUnitary`, which is beyond current scope; should be included in a future story of implementing sparse for generic controlled ops.
 - Any other hidden interfaces that were inherited from parent classes: same as above; require further discussions and extra stories/epics if needed. Current interfaces already provide enough features for `QubitUnitary` to be called in the pipeline. For example, `prod`, `sum` etc.

P.S. please note that the multiplication between operators and states is not yet involved here since it is a matter of `apply_operation` which should be resolved in some follow-ups (e.g. https://github.com/PennyLaneAI/pennylane/pull/6883#issuecomment-2613345018, or some other refined smaller PR).

**Related GitHub Issues:**
[sc-82068]

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/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.
JerryChen97: Trigger the CI to check what's missing for coverage. Still WIP
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6889?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
Attention: Patch coverage is `84.61538%` with `4 lines` in your changes missing coverage. Please review.
> Project coverage is 99.58%. Comparing base [(`8cefc9a`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/8cefc9a7d549b0a1644ff1f719d8f3c6fcf9bf8d?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`e19ace6`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/e19ace6726c78ef709e53d55b93457b22ea5a1fc?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 1 commits behind head on master.

| [Files with missing lines](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6889?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Patch % | Lines |
|---|---|---|
| [pennylane/math/multi\_dispatch.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6889?src=pr&el=tree&filepath=pennylane%2Fmath%2Fmulti_dispatch.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvbXVsdGlfZGlzcGF0Y2gucHk=) | 66.66% | [2 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6889?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) |
| [pennylane/ops/qubit/matrix\_ops.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6889?src=pr&el=tree&filepath=pennylane%2Fops%2Fqubit%2Fmatrix_ops.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9xdWJpdC9tYXRyaXhfb3BzLnB5) | 90.00% | [2 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6889?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #6889      +/-   ##
==========================================
- Coverage   99.59%   99.58%   -0.01%     
==========================================
  Files         480      480              
  Lines       45474    45499      +25     
==========================================
+ Hits        45289    45310      +21     
- Misses        185      189       +4     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6889?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

JerryChen97: > Nice work so far. @JerryChen97, can you think of a way to accept any scipy sparse representation, and internally convert it to CSR? If you do that in a helper function it will be easier to use it everywhere. [This](https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.issparse.html#scipy.sparse.issparse) might help you do the trick.

@AmintorDusko My plan is to put `issparse` in the init of both `StatePrep` and `QubitUnitary`, and then convert to `csr_matrix` specifically. 
It might be also helpful if we know whether or not this is necessary. My own experience based on past practice is `SciPy.sparse` works fastest with `csr` or `csc` (e.g. one can try freely `scipy.sparse.kron` over 20 sites without enforcing everyone to be `csr` and it will be super slow) but not sure if this is still the generic case right now.
AmintorDusko: > > Nice work so far. @JerryChen97, can you think of a way to accept any scipy sparse representation, and internally convert it to CSR? If you do that in a helper function it will be easier to use it everywhere. [This](https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.issparse.html#scipy.sparse.issparse) might help you do the trick.
> 
> @AmintorDusko My plan is to put `issparse` in the init of both `StatePrep` and `QubitUnitary`, and then convert to `csr_matrix` specifically. It might be also helpful if we know whether or not this is necessary. My own experience based on past practice is `SciPy.sparse` works fastest with `csr` or `csc` (e.g. one can try freely `scipy.sparse.kron` over 20 sites without enforcing everyone to be `csr` and it will be super slow) but not sure if this is still the generic case right now.

If I understood what you are saying correctly, this is exactly what I proposed to you.
JerryChen97: > This seems flawless to me! 🚀
> 
> Just one question. Except for the `_controlled` method and any other interfaces inherited from parent classes (as well as multiplication between operators and states), is there something else that is supported without sparse matrices but not supported with sparse matrices?

I realized that this need might be connected with the other request that we keep the whole process of decompositions of QubitUnitary sparse all the way along, both of which could be implemented by another story that fix the dispatches of `qml.math` for `scipy.sparse`. Let me explore if there's quick fix on this direction. If it works I might ask you guys to review another one PR on it.
AmintorDusko: > @AmintorDusko @PietropaoloFrisoni To accommodate our requirements for preservation of the sparsity, I've made quite some changes inside our `qml.math` module as well as some necessary touch on 1-qb and 2-qb decompositions. Besides, `RX`, `RY` and `RZ` have been manually added `sparse_matrix` compute logic. Also, `issparse` checks have been put in the beginning of `QubitUnitary` as well as `StatePrep`, and we use scipy builtin `tocsr` method to enforce everyone to be in the format of CSR. However, a helper function to implement `det` of sparse unitary matrices made use of `csc_matrix`, because no matter how `scipy.sparse.linalg.splu` will always convert things into CSC, or raise a warning, therefore we just convert it in PL to avoid warnings.

Converting internally to however is most efficient is not a problem. SciPy converts to CSC because this is a column-based operation. Internally we can make any sensible conversions, we only need consistency (especially for the output). The **CSR** format is a good general-purpose format and default choice. It provides efficient row-based access, fast matrix-vector products, good performance for arithmetic operations and compact storage. On the other hand, slow for column access, and inserting/deleting elements. The **CSC** is similar to CSR but is optimized for column-based operations. This will make it faster for matrix-vector products (especially when multiplying the matrix on the left: A @ v and some similar algebra operations. As before, this format will be slower for row access and inserting/deleting elements. Another usual format is the **COO**, this one is great for building sparse matrices, and adding/changing entries. This format will be slower and more memory-intensive everywhere else. SciPy also provides other interesting formats each with their specificities. 
JerryChen97: > > @AmintorDusko @PietropaoloFrisoni To accommodate our requirements for preservation of the sparsity, I've made quite some changes inside our `qml.math` module as well as some necessary touch on 1-qb and 2-qb decompositions. Besides, `RX`, `RY` and `RZ` have been manually added `sparse_matrix` compute logic. Also, `issparse` checks have been put in the beginning of `QubitUnitary` as well as `StatePrep`, and we use scipy builtin `tocsr` method to enforce everyone to be in the format of CSR. However, a helper function to implement `det` of sparse unitary matrices made use of `csc_matrix`, because no matter how `scipy.sparse.linalg.splu` will always convert things into CSC, or raise a warning, therefore we just convert it in PL to avoid warnings.
> 
> Converting internally to however is most efficient is not a problem. SciPy converts to CSC because this is a column-based operation. Internally we can make any sensible conversions, we only need consistency (especially for the output). The **CSR** format is a good general-purpose format and default choice. It provides efficient row-based access, fast matrix-vector products, good performance for arithmetic operations and compact storage. On the other hand, slow for column access, and inserting/deleting elements. The **CSC** is similar to CSR but is optimized for column-based operations. This will make it faster for matrix-vector products (especially when multiplying the matrix on the left: A @ v and some similar algebra operations. As before, this format will be slower for row access and inserting/deleting elements. Another usual format is the **COO**, this one is great for building sparse matrices, and adding/changing entries. This format will be slower and more memory-intensive everywhere else. SciPy also provides other interesting formats each with their specificities.

Thanks a lot for posting the more detailed information here! This will be an important reference when in the future people look back on our choice and implementation.
BTW, on the matrix operations employed in PL, I noticed that we default the batched dimension on the first dimension, so eventually it might be convenient if we use `v @ A.T` to update the states `apply_operation`. This is also what I've been trying to use in the next PR https://github.com/PennyLaneAI/pennylane/pull/6883.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 585aab458..ee1bcfcd5 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -10,6 +10,31 @@
 
 <h3>Improvements 🛠</h3>
 
+* `qml.QubitUnitary` now accepts sparse CSR matrices (from `scipy.sparse`). This allows efficient representation of large unitaries with mostly zero entries. Note that sparse unitaries are still in early development and may not support all features of their dense counterparts.
+  [(#6889)](https://github.com/PennyLaneAI/pennylane/pull/6889)
+
+  ```pycon
+  >>> import numpy as np
+  >>> import pennylane as qml
+  >>> import scipy as sp
+  >>> U_dense = np.eye(4)  # 2-wire identity
+  >>> U_sparse = sp.sparse.csr_matrix(U_dense)
+  >>> op = qml.QubitUnitary(U_sparse, wires=[0, 1])
+  >>> print(op.matrix())
+  <Compressed Sparse Row sparse matrix of dtype 'float64'
+          with 4 stored elements and shape (4, 4)>
+    Coords        Values
+    (0, 0)        1.0
+    (1, 1)        1.0
+    (2, 2)        1.0
+    (3, 3)        1.0
+  >>> op.matrix().toarray()
+  array([[1., 0., 0., 0.],
+        [0., 1., 0., 0.],
+        [0., 0., 1., 0.],
+        [0., 0., 0., 1.]])
+  ```
+
 * Add a decomposition for multi-controlled global phases into a one-less-controlled phase shift.
   [(#6936)](https://github.com/PennyLaneAI/pennylane/pull/6936)
  
@@ -86,7 +111,7 @@
   >>> print(new_circuit.diff_method)
   'parameter-shift'
   ```
-  
+
 * Devices can now configure whether or not ML framework data is sent to them
   via an `ExecutionConfig.convert_to_numpy` parameter. End-to-end jitting on
   `default.qubit` is used if the user specified a `jax.random.PRNGKey` as a seed.
diff --git a/pennylane/math/single_dispatch.py b/pennylane/math/single_dispatch.py
index 6d9902d41..0598f466f 100644
--- a/pennylane/math/single_dispatch.py
+++ b/pennylane/math/single_dispatch.py
@@ -19,8 +19,10 @@ from importlib import import_module
 # pylint: disable=wrong-import-order
 import autoray as ar
 import numpy as np
+import scipy as sp
 from packaging.version import Version
 from scipy.linalg import block_diag as _scipy_block_diag
+from scipy.sparse.linalg import splu
 
 from .interface_utils import get_deep_interface
 from .utils import is_abstract
@@ -68,6 +70,67 @@ ar.register_function("scipy", "transpose", np.transpose)
 ar.register_function("scipy", "ndim", np.ndim)
 
 
+# -------------------------------- SciPy Sparse --------------------------------- #
+# the following is required to ensure that general SciPy sparse matrices are
+# not automatically 'unwrapped' to dense NumPy arrays. Note that we assume
+# that whenever the backend is 'scipy', the input is a SciPy sparse matrix.
+
+
+def _det_sparse(x):
+    """Compute determinant of sparse matrices without densification"""
+
+    assert sp.sparse.issparse(x), TypeError(f"Expected SciPy sparse, got {type(x)}")
+
+    x = sp.sparse.csr_matrix(x)
+    if x.shape == (2, 2):
+        # Direct array access
+        indptr, indices, data = x.indptr, x.indices, x.data
+        values = {(i, j): 0.0 for i in range(2) for j in range(2)}
+        for i in range(2):
+            for j_idx in range(indptr[i], indptr[i + 1]):
+                j = indices[j_idx]
+                values[(i, j)] = data[j_idx]
+        return values[(0, 0)] * values[(1, 1)] - values[(0, 1)] * values[(1, 0)]
+    return _generic_sparse_det(x)
+
+
+def _generic_sparse_det(A):
+    """Compute the determinant of a sparse matrix using LU decomposition."""
+
+    assert hasattr(A, "tocsc"), TypeError(f"Expected SciPy sparse, got {type(A)}")
+
+    A_csc = A.tocsc()
+    lu = splu(A_csc)
+    U_diag = lu.U.diagonal()
+    det_A = np.prod(U_diag)
+    parity = _permutation_parity(lu.perm_r)
+    return det_A * parity
+
+
+def _permutation_parity(perm):
+    """Compute the parity of a permutation."""
+
+    parity = 1
+    visited = [False] * len(perm)
+    for i in range(len(perm)):
+        if not visited[i]:
+            cycle_length = 0
+            j = i
+            while not visited[j]:
+                visited[j] = True
+                j = perm[j]
+                cycle_length += 1
+
+            if cycle_length:
+
+                parity *= (-1) ** (cycle_length - 1)
+    return parity
+
+
+ar.register_function("scipy", "linalg.det", _det_sparse)
+ar.register_function("scipy", "linalg.eigs", sp.sparse.linalg.eigs)
+ar.register_function("scipy", "trace", lambda x: x.trace())
+
 # -------------------------------- NumPy --------------------------------- #
 
 ar.register_function("numpy", "flatten", lambda x: x.flatten())
diff --git a/pennylane/math/utils.py b/pennylane/math/utils.py
index 823c561bc..0c333d58d 100644
--- a/pennylane/math/utils.py
+++ b/pennylane/math/utils.py
@@ -16,6 +16,7 @@
 # pylint: disable=wrong-import-order
 import autoray as ar
 import numpy as _np
+import scipy as sp
 
 # pylint: disable=import-outside-toplevel
 from autograd.numpy.numpy_boxes import ArrayBox
@@ -174,6 +175,9 @@ def convert_like(tensor1, tensor2):
         dev = tensor2.device
         return np.asarray(tensor1, device=dev, like=interface)
 
+    if interface == "scipy":
+        return sp.sparse.csr_matrix(tensor1)
+
     return np.asarray(tensor1, like=interface)
 
 
diff --git a/pennylane/ops/op_math/decompositions/single_qubit_unitary.py b/pennylane/ops/op_math/decompositions/single_qubit_unitary.py
index eab40d549..d8ef8be90 100644
--- a/pennylane/ops/op_math/decompositions/single_qubit_unitary.py
+++ b/pennylane/ops/op_math/decompositions/single_qubit_unitary.py
@@ -14,8 +14,10 @@
 """Contains transforms and helpers functions for decomposing arbitrary unitary
 operations into elementary gates.
 """
+from functools import singledispatch
 
 import numpy as np
+import scipy as sp
 
 import pennylane as qml
 from pennylane import math
@@ -42,11 +44,17 @@ def _convert_to_su2(U, return_global_phase=False):
     with np.errstate(divide="ignore", invalid="ignore"):
         determinants = math.linalg.det(U)
     phase = math.angle(determinants) / 2
-    U = math.cast_like(U, determinants) * math.exp(-1j * math.cast_like(phase, 1j))[:, None, None]
+    U = (
+        U * math.exp(-1j * phase)
+        if sp.sparse.issparse(U)
+        else math.cast_like(U, determinants)
+        * math.exp(-1j * math.cast_like(phase, 1j))[:, None, None]
+    )
 
     return (U, phase) if return_global_phase else U
 
 
+@singledispatch
 def _zyz_get_rotation_angles(U):
     r"""Computes the rotation angles :math:`\phi`, :math:`\theta`, :math:`\omega`
     for a unitary :math:`U` that is :math:`SU(2)`
@@ -91,6 +99,49 @@ def _zyz_get_rotation_angles(U):
     return phis, thetas, omegas
 
 
+@_zyz_get_rotation_angles.register(sp.sparse.csr_matrix)
+def _zyz_get_rotation_angles_sparse(U):
+    r"""Computes the rotation angles :math:`\phi`, :math:`\theta`, :math:`\omega`
+    for a unitary :math:`U` that is :math:`SU(2)`, sparse case
+
+    Args:
+        U (array[complex]): A matrix that is :math:`SU(2)`
+
+    Returns:
+        tuple[array[float]]: A tuple containing the rotation angles
+            :math:`\phi`, :math:`\theta`, :math:`\omega`
+
+    """
+
+    assert sp.sparse.issparse(U), "Do not use this method if U is not sparse"
+
+    u00 = U[0, 0]
+    u01 = U[0, 1]
+    u10 = U[1, 0]
+
+    # For batched U or single U with non-zero off-diagonal, compute the
+    # normal decomposition instead
+    off_diagonal_elements = math.clip(math.abs(u01), 0, 1)
+    thetas = 2 * math.arcsin(off_diagonal_elements)
+
+    # Compute phi and omega from the angles of the top row; use atan2 to keep
+    # the angle within -np.pi and np.pi
+    angles_U00 = math.arctan2(math.imag(u00), math.real(u00))
+    angles_U10 = math.arctan2(math.imag(u10), math.real(u10))
+
+    phis = -angles_U10 - angles_U00
+    omegas = angles_U10 - angles_U00
+
+    phis, thetas, omegas = map(math.squeeze, [phis, thetas, omegas])
+
+    # Normalize the angles
+    phis = phis % (4 * np.pi)
+    thetas = thetas % (4 * np.pi)
+    omegas = omegas % (4 * np.pi)
+
+    return phis, thetas, omegas
+
+
 def _rot_decomposition(U, wire, return_global_phase=False):
     r"""Compute the decomposition of a single-qubit matrix :math:`U` in terms of
     elementary operations, as a single :class:`.RZ` gate or a :class:`.Rot` gate.
@@ -167,7 +218,7 @@ def _get_single_qubit_rot_angles_via_matrix(
     of the matrix of the target operation using ZYZ rotations.
     """
     # Cast to batched format for more consistent code
-    U = math.expand_dims(U, axis=0) if len(U.shape) == 2 else U
+    U = math.expand_dims(U, axis=0) if len(U.shape) == 2 and not sp.sparse.issparse(U) else U
 
     # Convert to SU(2) format and extract global phase
     U_su2, global_phase = _convert_to_su2(U, return_global_phase=True)
diff --git a/pennylane/ops/op_math/decompositions/two_qubit_unitary.py b/pennylane/ops/op_math/decompositions/two_qubit_unitary.py
index de223d5f3..666a592ba 100644
--- a/pennylane/ops/op_math/decompositions/two_qubit_unitary.py
+++ b/pennylane/ops/op_math/decompositions/two_qubit_unitary.py
@@ -17,6 +17,7 @@ unitary operations into elementary gates.
 import warnings
 
 import numpy as np
+import scipy as sp
 
 import pennylane as qml
 from pennylane import math
@@ -102,6 +103,9 @@ q_one_cnot = (1 / np.sqrt(2)) * np.array(
 )
 
 
+global_arrays_name = ["E", "Edag", "CNOT01", "CNOT10", "SWAP", "S_SX", "v_one_cnot", "q_one_cnot"]
+
+
 def _convert_to_su4(U):
     r"""Convert a 4x4 matrix to :math:`SU(4)`.
 
@@ -621,6 +625,11 @@ def two_qubit_decomposition(U, wires):
     _check_differentiability_warning(U)
     # First, we note that this method works only for SU(4) gates, meaning that
     # we need to rescale the matrix by its determinant.
+    if sp.sparse.issparse(U):
+        # Convert all the global elements to sparse matrices in-place
+        for name in global_arrays_name:
+            array = globals()[name]
+            globals()[name] = sp.sparse.csr_matrix(array)
     U = _convert_to_su4(U)
 
     # The next thing we will do is compute the number of CNOTs needed, as this affects
diff --git a/pennylane/ops/qubit/matrix_ops.py b/pennylane/ops/qubit/matrix_ops.py
index 2ca5bd9d5..0a79dad6d 100644
--- a/pennylane/ops/qubit/matrix_ops.py
+++ b/pennylane/ops/qubit/matrix_ops.py
@@ -21,7 +21,9 @@ from itertools import product
 from typing import Optional, Union
 
 import numpy as np
+import scipy as sp
 from scipy.linalg import fractional_matrix_power
+from scipy.sparse import csr_matrix
 
 import pennylane as qml
 from pennylane import numpy as pnp
@@ -78,6 +80,10 @@ class QubitUnitary(Operation):
     r"""QubitUnitary(U, wires)
     Apply an arbitrary unitary matrix with a dimension that is a power of two.
 
+    .. warning::
+
+        The sparse matrix representation of QubitUnitary is still under development. Currently we only support a limited set of interfaces that preserve the sparsity of the matrix, including ..method::`adjoint`, ..method::`pow`, ..method::`compute_sparse_matrix` and ..method::`compute_decomposition`. Differentiability is not supported for sparse matrices.
+
     **Details:**
 
     * Number of wires: Any (the operation can act on any number of wires)
@@ -86,7 +92,7 @@ class QubitUnitary(Operation):
     * Gradient recipe: None
 
     Args:
-        U (array[complex]): square unitary matrix
+        U (array[complex] or csr_matrix): square unitary matrix
         wires (Sequence[int] or int): the wire(s) the operation acts on
         id (str): custom label given to an operator instance,
             can be useful for some applications where the instance has to be identified
@@ -121,7 +127,7 @@ class QubitUnitary(Operation):
 
     def __init__(
         self,
-        U: TensorLike,
+        U: Union[TensorLike, csr_matrix],
         wires: WiresLike,
         id: Optional[str] = None,
         unitary_check: bool = False,
@@ -135,19 +141,16 @@ class QubitUnitary(Operation):
         if len(U_shape) not in {2, 3} or U_shape[-2:] != (dim, dim):
             raise ValueError(
                 f"Input unitary must be of shape {(dim, dim)} or (batch_size, {dim}, {dim}) "
-                f"to act on {len(wires)} wires."
+                f"to act on {len(wires)} wires. Got shape {U_shape} instead."
             )
 
+        # If the matrix is sparse, we need to convert it to a csr_matrix
+        if sp.sparse.issparse(U):
+            U = U.tocsr()
+
         # Check for unitarity; due to variable precision across the different ML frameworks,
         # here we issue a warning to check the operation, instead of raising an error outright.
-        if unitary_check and not (
-            qml.math.is_abstract(U)
-            or qml.math.allclose(
-                qml.math.einsum("...ij,...kj->...ik", U, qml.math.conj(U)),
-                qml.math.eye(dim),
-                atol=1e-6,
-            )
-        ):
+        if unitary_check and not self._unitary_check(U, dim):
             warnings.warn(
                 f"Operator {U}\n may not be unitary. "
                 "Verify unitarity of operation, or use a datatype with increased precision.",
@@ -156,6 +159,18 @@ class QubitUnitary(Operation):
 
         super().__init__(U, wires=wires, id=id)
 
+    @staticmethod
+    def _unitary_check(U, dim):
+        if isinstance(U, csr_matrix):
+            U_dagger = U.conjugate().transpose()
+            identity = sp.sparse.eye(dim, format="csr")
+            return sp.sparse.linalg.norm(U @ U_dagger - identity) < 1e-10
+        return qml.math.allclose(
+            qml.math.einsum("...ij,...kj->...ik", U, qml.math.conj(U)),
+            qml.math.eye(dim),
+            atol=1e-6,
+        )
+
     @staticmethod
     def compute_matrix(U: TensorLike):  # pylint: disable=arguments-differ
         r"""Representation of the operator as a canonical matrix in the computational basis (static method).
@@ -180,6 +195,26 @@ class QubitUnitary(Operation):
         """
         return U
 
+    @staticmethod
+    def compute_sparse_matrix(U: TensorLike):  # pylint: disable=arguments-differ
+        r"""Representation of the operator as a sparse matrix.
+
+        Args:
+            U (tensor_like): unitary matrix
+
+        Returns:
+            csr_matrix: sparse matrix representation
+
+        **Example**
+
+        >>> U = np.array([[0.98877108+0.j, 0.-0.14943813j], [0.-0.14943813j, 0.98877108+0.j]])
+        >>> qml.QubitUnitary.compute_sparse_matrix(U)
+        <2x2 sparse matrix of type '<class 'numpy.complex128'>'
+            with 2 stored elements in Compressed Sparse Row format>
+        """
+        U = qml.math.asarray(U, like="numpy")
+        return sp.sparse.csr_matrix(U)
+
     @staticmethod
     def compute_decomposition(U: TensorLike, wires: WiresLike):
         r"""Representation of the operator as a product of other operators (static method).
@@ -236,10 +271,17 @@ class QubitUnitary(Operation):
 
     def adjoint(self) -> "QubitUnitary":
         U = self.matrix()
+        if isinstance(U, csr_matrix):
+            adjoint_sp_mat = U.conjugate().transpose()
+            # Note: it is necessary to explicitly cast back to csr, or it will be come csc
+            return QubitUnitary(csr_matrix(adjoint_sp_mat), wires=self.wires)
         return QubitUnitary(qml.math.moveaxis(qml.math.conj(U), -2, -1), wires=self.wires)
 
     def pow(self, z: Union[int, float]):
         mat = self.matrix()
+        if isinstance(mat, csr_matrix):
+            pow_mat = sp.sparse.linalg.matrix_power(mat, z)
+            return [QubitUnitary(pow_mat, wires=self.wires)]
         if isinstance(z, int) and qml.math.get_deep_interface(mat) != "tensorflow":
             pow_mat = qml.math.linalg.matrix_power(mat, z)
         elif self.batch_size is not None or qml.math.shape(z) != ():
diff --git a/pennylane/ops/qubit/parametric_ops_single_qubit.py b/pennylane/ops/qubit/parametric_ops_single_qubit.py
index c0c8c1681..456a8c480 100644
--- a/pennylane/ops/qubit/parametric_ops_single_qubit.py
+++ b/pennylane/ops/qubit/parametric_ops_single_qubit.py
@@ -21,6 +21,7 @@ import functools
 from typing import Optional, Union
 
 import numpy as np
+import scipy as sp
 
 import pennylane as qml
 from pennylane.operation import Operation
@@ -113,6 +114,15 @@ class RX(Operation):
         js = -1j * s
         return qml.math.stack([stack_last([c, js]), stack_last([js, c])], axis=-2)
 
+    @staticmethod
+    def compute_sparse_matrix(theta):
+        return sp.sparse.csr_matrix(
+            [
+                [np.cos(theta / 2), -1j * np.sin(theta / 2)],
+                [-1j * np.sin(theta / 2), np.cos(theta / 2)],
+            ]
+        )
+
     def adjoint(self) -> "RX":
         return RX(-self.data[0], wires=self.wires)
 
@@ -209,6 +219,12 @@ class RY(Operation):
         s = (1 + 0j) * s
         return qml.math.stack([stack_last([c, -s]), stack_last([s, c])], axis=-2)
 
+    @staticmethod
+    def compute_sparse_matrix(theta):
+        return sp.sparse.csr_matrix(
+            [[np.cos(theta / 2), -np.sin(theta / 2)], [np.sin(theta / 2), np.cos(theta / 2)]]
+        )
+
     def adjoint(self) -> "RY":
         return RY(-self.data[0], wires=self.wires)
 
@@ -307,6 +323,10 @@ class RZ(Operation):
         diags = qml.math.exp(qml.math.outer(arg, signs))
         return diags[:, :, np.newaxis] * qml.math.cast_like(qml.math.eye(2, like=diags), diags)
 
+    @staticmethod
+    def compute_sparse_matrix(theta):
+        return sp.sparse.csr_matrix([[np.exp(-1j * theta / 2), 0], [0, np.exp(1j * theta / 2)]])
+
     @staticmethod
     def compute_eigvals(theta: TensorLike) -> TensorLike:  # pylint: disable=arguments-differ
         r"""Eigenvalues of the operator in the computational basis (static method).
diff --git a/pennylane/ops/qubit/state_preparation.py b/pennylane/ops/qubit/state_preparation.py
index 342267d79..ee6b57fe2 100644
--- a/pennylane/ops/qubit/state_preparation.py
+++ b/pennylane/ops/qubit/state_preparation.py
@@ -326,7 +326,8 @@ class StatePrep(StatePrepBase):
         validate_norm: bool = True,
     ):
         self.is_sparse = False
-        if isinstance(state, csr_matrix):
+        if sp.sparse.issparse(state):
+            state = state.tocsr()
             state = self._preprocess_csr(state, wires, None, normalize, validate_norm)
             self.is_sparse = True
         else:


Test patch:
diff --git a/tests/ops/qubit/test_matrix_ops.py b/tests/ops/qubit/test_matrix_ops.py
index 0348bf734..dcaa0a486 100644
--- a/tests/ops/qubit/test_matrix_ops.py
+++ b/tests/ops/qubit/test_matrix_ops.py
@@ -20,6 +20,7 @@ from functools import reduce
 import numpy as np
 import pytest
 from gate_data import H, I, S, T, X, Z
+from scipy.sparse import coo_matrix, csc_matrix, csr_matrix
 
 import pennylane as qml
 from pennylane import numpy as pnp
@@ -28,6 +29,137 @@ from pennylane.ops.qubit.matrix_ops import _walsh_hadamard_transform, fractional
 from pennylane.wires import Wires
 
 
+class TestQubitUnitaryCSR:
+    """Tests for using csr_matrix in QubitUnitary."""
+
+    def test_compute_sparse_matrix(self):
+        """Test that the compute_sparse_matrix method works correctly."""
+        U = np.array([[0, 1], [1, 0]])
+        op = qml.QubitUnitary.compute_sparse_matrix(U)
+        assert isinstance(op, csr_matrix)
+        assert np.allclose(op.toarray(), U)
+
+    def test_generic_sparse_convert_to_csr(self):
+        """Test that other generic sparse matrices can be converted to csr_matrix."""
+        # 4x4 Identity as a csr_matrix
+        dense = np.eye(4)
+        sparse = coo_matrix(dense)
+        op = qml.QubitUnitary(sparse, wires=[0, 1])
+        assert isinstance(op.matrix(), csr_matrix)
+        sparse = csc_matrix(dense)
+        op = qml.QubitUnitary(sparse, wires=[0, 1])
+        assert isinstance(op.matrix(), csr_matrix)
+
+    @pytest.mark.parametrize(
+        "dense",
+        [H, I, S, T, X, Z],
+    )
+    def test_csr_matrix_init_success(self, dense):
+        """Test that a valid 2-wire csr_matrix can be instantiated, covering necessary single-qubit gates."""
+        # 4x4 Identity as a csr_matrix
+        sparse = csr_matrix(dense)
+        op = qml.QubitUnitary(sparse, wires=[0])
+        assert isinstance(op.matrix(), csr_matrix)  # Should still be sparse
+        assert qml.math.allclose(op.matrix().toarray(), dense)
+
+    def test_csr_matrix_shape_mismatch(self):
+        """Test that shape mismatch with csr_matrix raises an error."""
+        dense = np.eye(2)  # Only 2x2
+        sparse = csr_matrix(dense)
+        with pytest.raises(ValueError, match="Input unitary must be of shape"):
+            qml.QubitUnitary(sparse, wires=[0, 1])
+
+    def test_csr_matrix_unitary_check_fail(self):
+        """Test that unitary_check warns if the matrix may not be unitary."""
+        dense = np.array([[1, 0], [0, 0.5]])  # Not a unitary
+        sparse = csr_matrix(dense)
+        with pytest.warns(UserWarning, match="may not be unitary"):
+            qml.QubitUnitary(sparse, wires=0, unitary_check=True)
+
+    def test_csr_matrix_pow_integer(self):
+        """Test that QubitUnitary.pow() works for integer exponents with csr_matrix."""
+        dense = np.eye(4)
+        sparse = csr_matrix(dense)
+        op = qml.QubitUnitary(sparse, wires=[0, 1])
+        powered_ops = op.pow(2)
+        assert len(powered_ops) == 1
+
+        powered_op = powered_ops[0]
+        assert isinstance(powered_op, qml.QubitUnitary)
+        assert isinstance(powered_op.matrix(), csr_matrix)
+        # The resulting matrix should still be the identity
+        final_mat = powered_op.matrix()
+        # If it's still sparse, compare .toarray()
+        if isinstance(final_mat, csr_matrix):
+            final_mat = final_mat.toarray()
+        assert qml.math.allclose(final_mat, dense)
+
+    @pytest.mark.parametrize(
+        "dense",
+        [
+            np.eye(4),
+            np.array(
+                [[1, 0, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0], [0, 1, 0, 0]]
+            ),  # sample permutation matrix
+        ],
+    )
+    def test_csr_matrix_adjoint(self, dense):
+        """Test that QubitUnitary.adjoint() works with csr_matrix, matching dense result."""
+        sparse = csr_matrix(dense)
+        op = qml.QubitUnitary(sparse, wires=[0, 1])
+        adj_op = op.adjoint()
+
+        assert isinstance(adj_op, qml.QubitUnitary)
+        assert isinstance(adj_op.matrix(), csr_matrix)
+
+        final_mat = adj_op.matrix()
+        # Compare with dense representation if still sparse
+        if isinstance(final_mat, csr_matrix):
+            final_mat = final_mat.toarray()
+
+        # For real/complex conjugate transpose, if dense is unitary, final_mat == dense^\dagger
+        expected = dense.conjugate().T
+        assert qml.math.allclose(final_mat, expected)
+
+    def test_csr_matrix_adjoint_large(self):
+        """Construct a large sparse matrix (e.g. 2^20 dimension) but only store minimal elements."""
+        N = 20
+        dim = 2**N
+
+        # For demonstration, let's just store a single 1 on the diagonal
+        row_indices = [12345]  # some arbitrary index < dim
+        col_indices = [12345]
+        data = [1.0]
+
+        sparse_large = csr_matrix((data, (row_indices, col_indices)), shape=(dim, dim))
+        with pytest.warns(UserWarning, match="may not be unitary"):
+            op = qml.QubitUnitary(sparse_large, wires=range(N), unitary_check=True)
+        adj_op = op.adjoint()
+
+        assert isinstance(adj_op, qml.QubitUnitary)
+        assert isinstance(adj_op.matrix(), csr_matrix)
+
+        # The single element should remain 1 at [12345,12345] after conjugate transpose
+        final_mat = adj_op.matrix()
+        assert final_mat[12345, 12345] == 1.0
+
+    def test_csr_matrix_decomposition(self):
+        """Test that QubitUnitary.decomposition() works with csr_matrix."""
+        # 4x4 Identity as a csr_matrix
+        dense = np.eye(4)
+        sparse = csr_matrix(dense)
+        op = qml.QubitUnitary(sparse, wires=[0, 1])
+        decomp = op.decomposition()
+        assert len(decomp) == 6
+
+        # 2x2 Identity as a csr_matrix
+        dense = np.eye(2)
+        sparse = csr_matrix(dense)
+        op = qml.QubitUnitary(sparse, wires=[0])
+        decomp = op.decomposition()
+        assert len(decomp) == 3
+
+
 class TestQubitUnitary:
     """Tests for the QubitUnitary class."""
 
diff --git a/tests/ops/qubit/test_state_prep.py b/tests/ops/qubit/test_state_prep.py
index 1c6975e64..49b299d06 100644
--- a/tests/ops/qubit/test_state_prep.py
+++ b/tests/ops/qubit/test_state_prep.py
@@ -413,6 +413,13 @@ class TestStateVector:
 class TestSparseStateVector:
     """Test the sparse_state_vector() method of various state-prep operations."""
 
+    def test_sparse_state_convert_to_csr(self):
+        """Test that the sparse_state_vector() method returns a csr_matrix."""
+        sp_vec = sp.sparse.coo_matrix([0, 0, 1, 0])
+        qsv_op = qml.StatePrep(sp_vec, wires=[0, 1])
+        ket = qsv_op.state_vector()
+        assert isinstance(ket, sp.sparse.csr_matrix)
+
     @pytest.mark.parametrize(
         "num_wires,wire_order,one_position",
         [