Repository: PennyLaneAI/pennylane
Pull request #7285 (2025-04-15 14:57:15)

Fix sparse `BlockEncode` bug
**Context:**
A cherrypicked PR from https://github.com/PennyLaneAI/pennylane/pull/7248/commits
`BlockEncode` initialized with sparse matrices seems not working properly when applied into states. We would like to match its strcuture similar with `QubitUnitary`

**Description of the Change:**
src:
 - Extract the logic of building encoder matrices and place them into a separate private helper method
 - Construct similar `issparse` logic following the same pattern as `QubitUnitary`

tests:
 - Added a unit test to make sure the fledged `BlockEncode` behaves as expected
 - Added an integration test to make sure it can be applied inside a qnode as the original issue expected

**Benefits:**
Users can use sparse types of `BlockEncode` as a normal operator to apply in their circuits now.

**Possible Drawbacks:**

**Related GitHub Issues:**
https://github.com/PennyLaneAI/pennylane/issues/7254
[sc-88769]


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/7285?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.64%. Comparing base [(`fa1a822`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/fa1a8228a39c8ee8a9e4bc96040ea74a3ba12eba?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`e5f817e`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/e5f817e2810d38322e2af4320dd5f8b40eaad44a?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.

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


```diff
@@           Coverage Diff           @@
##           master    #7285   +/-   ##
=======================================
  Coverage   99.64%   99.64%           
=======================================
  Files         513      513           
  Lines       50194    50210   +16     
=======================================
+ Hits        50018    50034   +16     
  Misses        176      176           
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7285?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).

<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>
JerryChen97: > Can we add some integration tests to make sure the bug is actually solved now?

Added here https://github.com/PennyLaneAI/pennylane/pull/7285/commits/1e9fe38207b8ae64579fbb250c42269166ff7683
JerryChen97: @astralcai 

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index f02ee8ba9..29f0ada13 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -155,6 +155,9 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* Fixes a bug where circuit execution fails with ``BlockEncode`` initialized with sparse matrices.
+  [(#7285)](https://github.com/PennyLaneAI/pennylane/pull/7285)
+
 * Adds an informative error if `qml.cond` is used with an abstract condition with
   jitting on `default.qubit` if capture is enabled.
   [(#7314)](https://github.com/PennyLaneAI/pennylane/pull/7314)
diff --git a/pennylane/ops/qubit/matrix_ops.py b/pennylane/ops/qubit/matrix_ops.py
index ec47748cb..7b4c7daf0 100644
--- a/pennylane/ops/qubit/matrix_ops.py
+++ b/pennylane/ops/qubit/matrix_ops.py
@@ -671,6 +671,20 @@ class BlockEncode(Operation):
         self.hyperparameters["norm"] = normalization
         self.hyperparameters["subspace"] = subspace
 
+        self._issparse = sp.sparse.issparse(A)
+
+    # pylint: disable=arguments-renamed, invalid-overridden-method
+    @property
+    def has_sparse_matrix(self) -> bool:
+        """bool: Whether the operator has a sparse matrix representation."""
+        return self._issparse
+
+    # pylint: disable=arguments-renamed, invalid-overridden-method
+    @property
+    def has_matrix(self) -> bool:
+        """bool: Whether the operator has a sparse matrix representation."""
+        return not self._issparse
+
     def _flatten(self) -> FlatPytree:
         return self.data, (self.wires, ())
 
@@ -704,45 +718,22 @@ class BlockEncode(Operation):
                [-0.07621992,  0.89117368, -0.2       , -0.4       ]])
         """
         A = params[0]
-        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
-                return qml.math.concat(lst, like=like, axis=axis)
-            return qml.math.hstack(lst) if h else qml.math.vstack(lst)
-
-        interface = qml.math.get_interface(A)
-
-        if qml.math.sum(shape_a) <= 2:
-            col1 = _stack([A, sqrt(1 - A * conj(A))], like=interface)
-            col2 = _stack([sqrt(1 - A * conj(A)), -conj(A)], like=interface)
-            u = _stack([col1, col2], h=True, like=interface)
-        else:
-            d1, d2 = shape_a
-            col1 = _stack(
-                [A, sqrtm(cast(eye(d2, like=A), A.dtype) - qml.math.transpose(conj(A)) @ A)],
-                like=interface,
-            )
-            col2 = _stack(
-                [
-                    sqrtm(cast(eye(d1, like=A), A.dtype) - A @ transpose(conj(A))),
-                    -transpose(conj(A)),
-                ],
-                like=interface,
+        subspace = hyperparams["subspace"]
+        if sp.sparse.issparse(A):
+            raise qml.operation.MatrixUndefinedError(
+                "The operator was initialized with a sparse matrix. Use sparse_matrix instead."
             )
-            u = _stack([col1, col2], h=True, like=interface)
+        return _process_blockencode(A, subspace)
 
-        if n + m < k:
-            r = k - (n + m)
-            col1 = _stack([u, zeros((r, n + m), like=A)], like=interface)
-            col2 = _stack([zeros((n + m, r), like=A), eye(r, like=A)], like=interface)
-            u = _stack([col1, col2], h=True, like=interface)
-
-        return u
+    @staticmethod
+    def compute_sparse_matrix(*params, **hyperparams):
+        A = params[0]
+        subspace = hyperparams["subspace"]
+        if sp.sparse.issparse(A):
+            return _process_blockencode(A, subspace)
+        raise qml.operation.SparseMatrixUndefinedError(
+            "The operator is initialized with a dense matrix, use the matrix method instead."
+        )
 
     def adjoint(self) -> "BlockEncode":
         A = self.parameters[0]
@@ -755,3 +746,48 @@ class BlockEncode(Operation):
         cache: Optional[dict] = None,
     ):
         return super().label(decimals=decimals, base_label=base_label or "BlockEncode", cache=cache)
+
+
+def _process_blockencode(A, subspace):
+    """
+    Process the BlockEncode operation.
+    """
+    n, m, k = 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
+            return qml.math.concat(lst, like=like, axis=axis)
+        return qml.math.hstack(lst) if h else qml.math.vstack(lst)
+
+    interface = qml.math.get_interface(A)
+
+    if qml.math.sum(shape_a) <= 2:
+        col1 = _stack([A, sqrt(1 - A * conj(A))], like=interface)
+        col2 = _stack([sqrt(1 - A * conj(A)), -conj(A)], like=interface)
+        u = _stack([col1, col2], h=True, like=interface)
+    else:
+        d1, d2 = shape_a
+        col1 = _stack(
+            [A, sqrtm(cast(eye(d2, like=A), A.dtype) - qml.math.transpose(conj(A)) @ A)],
+            like=interface,
+        )
+        col2 = _stack(
+            [
+                sqrtm(cast(eye(d1, like=A), A.dtype) - A @ transpose(conj(A))),
+                -transpose(conj(A)),
+            ],
+            like=interface,
+        )
+        u = _stack([col1, col2], h=True, like=interface)
+
+    if n + m < k:
+        r = k - (n + m)
+        col1 = _stack([u, zeros((r, n + m), like=A)], like=interface)
+        col2 = _stack([zeros((n + m, r), like=A), eye(r, like=A)], like=interface)
+        u = _stack([col1, col2], h=True, like=interface)
+
+    return u


Test patch:
diff --git a/tests/ops/qubit/test_matrix_ops.py b/tests/ops/qubit/test_matrix_ops.py
index 846b54f66..c8379033b 100644
--- a/tests/ops/qubit/test_matrix_ops.py
+++ b/tests/ops/qubit/test_matrix_ops.py
@@ -1005,6 +1005,48 @@ class TestUnitaryLabels:
 class TestBlockEncode:
     """Test the BlockEncode operation."""
 
+    @pytest.mark.parametrize(
+        ("input_matrix", "issparse"),
+        [
+            (np.array([[1, 0], [0, 1]]), False),
+            (X, False),
+            (csr_matrix([[1, 0], [0, 1]]), True),
+        ],
+    )
+    def test_property(self, input_matrix, issparse):
+        """Test that BlockEncode has the correct properties."""
+        op = qml.BlockEncode(input_matrix, wires=range(2))
+        assert op.has_matrix != issparse
+        assert op.has_sparse_matrix == issparse
+        if issparse:
+            # Test if there's correct error out
+            with pytest.raises(qml.operation.MatrixUndefinedError):
+                op.matrix()
+        else:
+            with pytest.raises(qml.operation.SparseMatrixUndefinedError):
+                op.sparse_matrix()
+
+    @pytest.mark.integration
+    @pytest.mark.parametrize(
+        "input_matrix",
+        [
+            np.array([[1, 0], [0, 1]]),
+            X,
+            csr_matrix([[1, 0], [0, 1]]),
+        ],
+    )
+    def test_applicable(self, input_matrix):
+        """Integration Test that BlockEncode can be applied onto a state."""
+        dev = qml.device("default.qubit")
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.BlockEncode(input_matrix, wires=range(2))
+            return qml.state()
+
+        state = circuit()
+        assert state.shape == (2**2,)
+
     @pytest.mark.parametrize(
         ("input_matrix", "wires", "expected_hyperparameters"),
         [
@@ -1390,7 +1432,7 @@ class TestBlockEncode:
         op = qml.BlockEncode(sparse_matrix, wires=range(num_wires))
 
         # Test the operator is unitary
-        mat = qml.matrix(op)
+        mat = op.sparse_matrix()
         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)