Repository: PennyLaneAI/pennylane
Pull request #6963 (2025-02-13 20:20:21)

`BlockEncode` accepts sparse matrices as input
**Context:**
As a following up of https://github.com/PennyLaneAI/pennylane/pull/6889, we make BlockEncode sparse

**Description of the Change:**
Overloaded many `qml.math` methods such that `BlockEncode` can automatically dispatch to sparse methods.

**Benefits:**
`BlockEncode` is able to accept and output sparse matrices now

**Possible Drawbacks:**
Current implementation is a bit fall-back, due to lack of scipy sparse support one many methods, e.g. `sum` over certain axis, `sqrt` as a matrix square root, etc

**Related GitHub Issues:**
[sc-82069]

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.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6963?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 99.59%. Comparing base [(`4444bac`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/4444bac8d57b6cff896ee1bd6b55fd15287c1f31?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`b046fdf`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/b046fdf925a4aff9b5b00129cc6accdabcec07c7?dropdown=coverage&el=desc&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    #6963   +/-   ##
=======================================
  Coverage   99.59%   99.59%           
=======================================
  Files         480      480           
  Lines       45695    45726   +31     
=======================================
+ Hits        45510    45541   +31     
  Misses        185      185           
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6963?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: Note: had a simple placeholder of sqrtm such that we can test the whole pipeline
JerryChen97: @andrijapau Reworkded the two `_allclose` sharing large chunk of logics here: https://github.com/PennyLaneAI/pennylane/pull/6963/commits/f43aac44874b96118bc527bea9eb80e37a48a1b1
I commented out previous implemnetations instead of direct deletion just incase we need to compare some details.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 03f58dcd3..a97d67375 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -53,6 +53,9 @@
 
 <h3>Improvements 🛠</h3>
 
+* `qml.BlockEncode` now accepts sparse input and outputs sparse matrices.
+  [(#6963)](https://github.com/PennyLaneAI/pennylane/pull/6963)
+
 * `Operator.sparse_matrix` now supports `format` parameter to specify the returned scipy sparse matrix format,
   with the default being `'csr'`
   [(#6995)](https://github.com/PennyLaneAI/pennylane/pull/6995)
diff --git a/pennylane/math/multi_dispatch.py b/pennylane/math/multi_dispatch.py
index b6728155e..9729ad78d 100644
--- a/pennylane/math/multi_dispatch.py
+++ b/pennylane/math/multi_dispatch.py
@@ -873,6 +873,9 @@ def norm(tensor, like=None, **kwargs):
     ):
         norm = _flat_autograd_norm
 
+    elif like == "scipy":
+        from scipy.sparse.linalg import norm
+
     else:
         from scipy.linalg import norm
 
diff --git a/pennylane/math/quantum.py b/pennylane/math/quantum.py
index 3750da1a6..8b9e29576 100644
--- a/pennylane/math/quantum.py
+++ b/pennylane/math/quantum.py
@@ -21,7 +21,7 @@ from string import ascii_letters as ABC
 import scipy as sp
 import scipy.sparse.linalg as spla
 from autoray import numpy as np
-from numpy import float64  # pylint:disable=wrong-import-order
+from numpy import float64, sqrt  # pylint:disable=wrong-import-order
 from scipy.sparse import csc_matrix, issparse
 
 import pennylane as qml
@@ -1033,6 +1033,8 @@ def _denman_beavers_iterations(mat, max_iter=100, tol=1e-13):
         LinAlgError: If matrix inversion fails
         ValueError: If NaN values or overflow are encountered during computation
     """
+    if mat.shape == (1, 1):
+        return sqrt(mat)
     try:
         mat = csc_matrix(mat)
         Y = mat
diff --git a/pennylane/math/single_dispatch.py b/pennylane/math/single_dispatch.py
index 383374eb1..73cef2fc1 100644
--- a/pennylane/math/single_dispatch.py
+++ b/pennylane/math/single_dispatch.py
@@ -65,6 +65,7 @@ ar.register_function("builtins", "logical_xor", lambda x, y: x ^ y)
 # qml.SparseHamiltonian are not automatically 'unwrapped' to dense NumPy arrays.
 ar.register_function("scipy", "to_numpy", lambda x: x)
 ar.register_function("scipy", "coerce", lambda x: x)
+ar.register_function("scipy", "array", lambda x: x)
 ar.register_function("scipy", "shape", np.shape)
 ar.register_function("scipy", "dot", np.dot)
 ar.register_function("scipy", "conj", np.conj)
@@ -191,6 +192,10 @@ ar.register_function("scipy", "reshape", lambda x, new_shape: x.reshape(new_shap
 ar.register_function("scipy", "real", lambda x: x.real)
 ar.register_function("scipy", "imag", lambda x: x.imag)
 ar.register_function("scipy", "size", lambda x: np.prod(x.shape))
+ar.register_function("scipy", "eye", sp.sparse.eye)
+ar.register_function("scipy", "zeros", sp.sparse.csr_matrix)
+ar.register_function("scipy", "hstack", sp.sparse.hstack)
+ar.register_function("scipy", "vstack", sp.sparse.vstack)
 
 # -------------------------------- NumPy --------------------------------- #
 
diff --git a/pennylane/math/utils.py b/pennylane/math/utils.py
index 0c333d58d..834731c1f 100644
--- a/pennylane/math/utils.py
+++ b/pennylane/math/utils.py
@@ -57,12 +57,88 @@ def allequal(tensor1, tensor2, **kwargs):
     return np.all(t1 == t2, **kwargs)
 
 
+def _allclose_sparse(a, b, rtol=1e-05, atol=1e-08):
+    """Compare two sparse matrices for approximate equality.
+
+    Args:
+        a, b: scipy sparse matrices to compare
+        rtol (float): relative tolerance
+        atol (float): absolute tolerance
+
+    Returns:
+        bool: True if matrices are approximately equal
+    """
+    if (a != b).nnz == 0:
+        return True
+
+    diff = abs(a - b)
+
+    # Handle cases where the matrix might be empty
+    max_diff = diff.data.max() if diff.nnz > 0 else 0
+    max_b = abs(b).data.max() if b.nnz > 0 else 0
+
+    return max_diff <= atol + rtol * max_b
+
+
+def _allclose_mixed(a, b, rtol=1e-05, atol=1e-08, b_is_sparse=True):
+    """Helper function for comparing dense and sparse matrices with correct tolerance reference.
+
+    Args:
+        a: first matrix (dense or sparse)
+        b: second matrix (sparse or dense)
+        rtol: relative tolerance
+        atol: absolute tolerance
+        b_is_sparse: True if b is sparse matrix, False if a is sparse matrix
+
+    Returns:
+        bool: True if matrices are approximately equal
+    """
+    sparse = b if b_is_sparse else a
+    dense = a if b_is_sparse else b
+
+    if sparse.nnz == 0:
+        return np.allclose(dense, 0, rtol=rtol, atol=atol)
+
+    if dense.shape != sparse.shape:
+        return False
+
+    SIZE_THRESHOLD = 10000
+    if np.prod(dense.shape) < SIZE_THRESHOLD:
+        # Use dense comparison but maintain b as reference
+        if b_is_sparse:
+            return np.allclose(a, sparse.toarray(), rtol=rtol, atol=atol)
+        return np.allclose(sparse.toarray(), b, rtol=rtol, atol=atol)
+
+    dense_coords = dense.nonzero()
+    sparse_coords = sparse.nonzero()
+
+    coord_diff = set(zip(*dense_coords)) ^ set(zip(*sparse_coords))
+    if coord_diff:
+        return False
+
+    # Maintain asymmetric comparison with correct reference
+    if b_is_sparse:
+        a_data = dense[dense_coords]
+        b_data = sparse.data
+    else:
+        a_data = sparse.data
+        b_data = dense[sparse_coords]
+    return np.allclose(a_data, b_data, rtol=rtol, atol=atol)
+
+
 def allclose(a, b, rtol=1e-05, atol=1e-08, **kwargs):
     """Wrapper around np.allclose, allowing tensors ``a`` and ``b``
     to differ in type"""
     try:
         # Some frameworks may provide their own allclose implementation.
         # Try and use it if available.
+        if sp.sparse.issparse(a) and sp.sparse.issparse(b):
+            return _allclose_sparse(a, b, rtol=rtol, atol=atol)
+        if sp.sparse.issparse(a):
+            # pylint: disable=arguments-out-of-order
+            return _allclose_mixed(a, b, rtol=rtol, atol=atol, b_is_sparse=False)
+        if sp.sparse.issparse(b):
+            return _allclose_mixed(a, b, rtol=rtol, atol=atol, b_is_sparse=True)
         res = np.allclose(a, b, rtol=rtol, atol=atol, **kwargs)
     except (TypeError, AttributeError, ImportError, RuntimeError):
         # Otherwise, convert the input to NumPy arrays.
diff --git a/pennylane/ops/qubit/matrix_ops.py b/pennylane/ops/qubit/matrix_ops.py
index 6e9338819..5351006fa 100644
--- a/pennylane/ops/qubit/matrix_ops.py
+++ b/pennylane/ops/qubit/matrix_ops.py
@@ -27,7 +27,17 @@ from scipy.sparse import csr_matrix
 
 import pennylane as qml
 from pennylane import numpy as pnp
-from pennylane.math import cast, conj, eye, norm, sqrt, sqrt_matrix, transpose, zeros
+from pennylane.math import (
+    cast,
+    conj,
+    eye,
+    norm,
+    sqrt,
+    sqrt_matrix,
+    sqrt_matrix_sparse,
+    transpose,
+    zeros,
+)
 from pennylane.operation import AnyWires, DecompositionUndefinedError, FlatPytree, Operation
 from pennylane.typing import TensorLike
 from pennylane.wires import Wires, WiresLike
@@ -665,6 +675,8 @@ class BlockEncode(Operation):
         n, m, k = hyperparams["subspace"]
         shape_a = qml.math.shape(A)
 
+        sqrtm = sqrt_matrix_sparse if sp.sparse.issparse(A) else sqrt_matrix
+
         def _stack(lst, h=False, like=None):
             if like == "tensorflow":
                 axis = 1 if h else 0
@@ -680,12 +692,12 @@ class BlockEncode(Operation):
         else:
             d1, d2 = shape_a
             col1 = _stack(
-                [A, sqrt_matrix(cast(eye(d2, like=A), A.dtype) - qml.math.transpose(conj(A)) @ A)],
+                [A, sqrtm(cast(eye(d2, like=A), A.dtype) - qml.math.transpose(conj(A)) @ A)],
                 like=interface,
             )
             col2 = _stack(
                 [
-                    sqrt_matrix(cast(eye(d1, like=A), A.dtype) - A @ transpose(conj(A))),
+                    sqrtm(cast(eye(d1, like=A), A.dtype) - A @ transpose(conj(A))),
                     -transpose(conj(A)),
                 ],
                 like=interface,


Test patch:
diff --git a/tests/math/test_functions.py b/tests/math/test_functions.py
index 8ce1bd701..f8afa50d6 100644
--- a/tests/math/test_functions.py
+++ b/tests/math/test_functions.py
@@ -19,6 +19,7 @@ from unittest.mock import patch
 
 import numpy as onp
 import pytest
+import scipy as sp
 from autograd.numpy.numpy_boxes import ArrayBox
 
 import pennylane as qml
@@ -254,6 +255,82 @@ def test_allclose(t1, t2):
     assert res == expected
 
 
+class TestAllCloseSparse:
+    """Test that the sparse-matrix specialized allclose functions works well"""
+
+    def test_dense_sparse_small_matrix(self):
+        """Test comparing small dense and sparse matrices"""
+        dense = np.array([[1, 0, 2], [0, 3, 0]])
+        sparse = sp.sparse.csr_matrix(dense)
+
+        assert fn.allclose(dense, sparse)
+        assert fn.allclose(sparse, dense)
+
+    def test_dense_sparse_zero_nnz(self):
+        """Test comparing dense matrix with empty sparse matrix"""
+        dense = np.zeros((2, 3))
+        sparse = sp.sparse.csr_matrix(dense)
+
+        assert fn.allclose(dense, sparse)
+        assert fn.allclose(sparse, dense)
+
+    def test_dense_sparse_different_shapes(self):
+        """Test comparing matrices with different shapes"""
+        dense = np.array([[1, 2], [3, 4]])
+        sparse = sp.sparse.csr_matrix(np.array([[1, 2]]))
+
+        assert not fn.allclose(dense, sparse)
+        assert not fn.allclose(sparse, dense)
+
+    def test_dense_sparse_large_matrix(self):
+        """Test comparing large dense and sparse matrices"""
+        n = 200
+        dense = np.eye(n)
+        sparse = sp.sparse.eye(n)
+
+        assert fn.allclose(dense, sparse)
+        assert fn.allclose(sparse, dense)
+
+        # When size is large enough, a very small perturbation
+        # will override the tolerance.
+        dense[-1, 0] = np.finfo(float).eps
+        assert not fn.allclose(dense, sparse)
+        assert not fn.allclose(sparse, dense)
+
+    def test_sparse_sparse_large_matrix(self):
+        """Test comparing large dense and sparse matrices"""
+        n = 200
+        dense = np.eye(n)
+        sparse0 = sp.sparse.csr_matrix(dense)
+        sparse = sp.sparse.eye(n)
+
+        assert fn.allclose(sparse0, sparse)
+
+        dense[-1, 0] = 0.001
+
+        sparse0 = sp.sparse.csr_matrix(dense)
+
+        assert not fn.allclose(sparse0, sparse)
+
+    def test_dense_sparse_different_nonzero(self):
+        """Test comparing matrices with different nonzero patterns"""
+        dense = np.array([[1, 0], [0, 1]])
+        sparse = sp.sparse.csr_matrix(np.array([[1, 1], [0, 1]]))
+
+        assert not fn.allclose(dense, sparse)
+        assert not fn.allclose(sparse, dense)
+
+    @pytest.mark.parametrize("rtol,atol", [(1e-7, 1e-8), (1e-5, 1e-6)])
+    def test_dense_sparse_tolerances(self, rtol, atol):
+        """Test comparing matrices with different tolerances"""
+        dense = np.array([[1.0, 0.0], [0.0, 1.0 + 1e-7]])
+        sparse = sp.sparse.csr_matrix(np.array([[1.0, 0.0], [0.0, 1.0]]))
+
+        allclose_result = fn.allclose(dense, sparse, rtol=rtol, atol=atol)
+        expected = np.allclose(dense, sparse.toarray(), rtol=rtol, atol=atol)
+        assert allclose_result == expected
+
+
 test_angle_data = [
     [1.0, 1.0j, 1 + 1j],
     [1.0, 1.0j, 1 + 1j],
diff --git a/tests/ops/qubit/test_matrix_ops.py b/tests/ops/qubit/test_matrix_ops.py
index cf5e0efaa..e8e5651a0 100644
--- a/tests/ops/qubit/test_matrix_ops.py
+++ b/tests/ops/qubit/test_matrix_ops.py
@@ -1290,6 +1290,56 @@ class TestBlockEncode:
 
         assert circuit(input_matrix) == output_value
 
+    @pytest.mark.parametrize(
+        "matrix_data",
+        [
+            {
+                "data": [0.1, 0.2, 0.3] * 4,
+                "indices": [0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3],
+                "indptr": [0, 3, 6, 9, 12],
+                "shape": (4, 8),
+            },
+            {
+                # A smaller example
+                "data": [1.0, 2.0, 3.0],
+                "indices": [0, 1, 2],
+                "indptr": [0, 3],
+                "shape": (1, 3),
+            },
+            {
+                # A small 2x2 square example
+                "data": [1.0, 2.0],
+                "indices": [0, 1],
+                "indptr": [0, 2, 2],
+                "shape": (2, 2),
+            },
+            {
+                # Another 3x3 square example
+                "data": [0.5, 1.1, 1.2, 0.3],
+                "indices": [0, 2, 1, 2],
+                "indptr": [0, 2, 3, 4],
+                "shape": (3, 3),
+            },
+        ],
+    )
+    @pytest.mark.parametrize("format", ["coo", "csr", "csc", "bsr"])
+    def test_sparse_matrix(self, matrix_data, format):
+        """Test that the BlockEncode works well with various sparse matrices."""
+        data = matrix_data["data"]
+        indices = matrix_data["indices"]
+        indptr = matrix_data["indptr"]
+        shape = matrix_data["shape"]
+
+        num_wires = 5
+        sparse_matrix = csr_matrix((data, indices, indptr), shape=shape).asformat(format)
+        op = qml.BlockEncode(sparse_matrix, wires=range(num_wires))
+
+        # Test the operator is unitary
+        mat = qml.matrix(op)
+        assert np.allclose(np.eye(mat.shape[0]), (mat @ mat.T.conj()).toarray())
+        mat_dense = qml.matrix(qml.BlockEncode(sparse_matrix.toarray(), wires=range(num_wires)))
+        assert qml.math.allclose(mat, mat_dense)
+
 
 class TestInterfaceMatricesLabel:
     """Test different interface matrices with qubit."""