Repository: PennyLaneAI/pennylane
Pull request #5110 (2024-01-25 21:20:02)

Make BlockEncode constructor JIT-friendly
The new BlockEncode operator is currently not JIT friendly due to its constructor (whereas the matrix method already seems to be).

Conditionals are replaced with equivalent mathematical functions.

[sc-54971]

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-commenter: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5110?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:
> Comparison is base [(`e69be08`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/e69be08529514ce6e5bb13383b5e3cbb20509c86?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.68% compared to head [(`61f7cb4`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5110?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.68%.
> Report is 3 commits behind head on master.


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


```diff
@@            Coverage Diff             @@
##           master    #5110      +/-   ##
==========================================
- Coverage   99.68%   99.68%   -0.01%     
==========================================
  Files         394      394              
  Lines       36323    36071     -252     
==========================================
- Hits        36208    35956     -252     
  Misses        115      115              
```



</details>

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


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 1d4a4b5e1..6d2415e63 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -186,10 +186,14 @@
 * Faster `qml.probs` measurements due to an optimization in `_samples_to_counts`.
   [(#5145)](https://github.com/PennyLaneAI/pennylane/pull/5145)
 
+* Ensure the `BlockEncode` operator is JIT-compatible with JAX.
+  [(#5110)](https://github.com/PennyLaneAI/pennylane/pull/5110)
+
 * Cuts down on performance bottlenecks in converting a `PauliSentence` to a `Sum`.
   [(#5141)](https://github.com/PennyLaneAI/pennylane/pull/5141)
   [(#5150)](https://github.com/PennyLaneAI/pennylane/pull/5150)
 
+
 * The `qml.qsvt` function uses `qml.GlobalPhase` instead of `qml.exp` to define global phase.
   [(#5105)](https://github.com/PennyLaneAI/pennylane/pull/5105)
 
@@ -487,6 +491,7 @@ Diksha Dhawan,
 Lillian Frederiksen,
 Eugenio Gigante,
 Diego Guala,
+David Ittah,
 Soran Jahangiri,
 Korbinian Kottmann,
 Christina Lee,
diff --git a/pennylane/ops/qubit/matrix_ops.py b/pennylane/ops/qubit/matrix_ops.py
index 7b9d5bc71..dace8fb04 100644
--- a/pennylane/ops/qubit/matrix_ops.py
+++ b/pennylane/ops/qubit/matrix_ops.py
@@ -515,22 +515,26 @@ class BlockEncode(Operation):
     """Gradient computation method."""
 
     def __init__(self, A, wires, id=None):
+        wires = Wires(wires)
         shape_a = qml.math.shape(A)
         if shape_a == () or all(x == 1 for x in shape_a):
             A = qml.math.reshape(A, [1, 1])
-
-        wires = Wires(wires)
-        if pnp.sum(shape_a) <= 2:
-            normalization = A if A > 1 else 1
+            normalization = qml.math.abs(A)
             subspace = (1, 1, 2 ** len(wires))
+
         else:
-            normalization = max(
+            if len(shape_a) == 1:
+                A = qml.math.reshape(A, [1, len(A)])
+                shape_a = qml.math.shape(A)
+
+            normalization = qml.math.maximum(
                 norm(A @ qml.math.transpose(qml.math.conj(A)), ord=pnp.inf),
                 norm(qml.math.transpose(qml.math.conj(A)) @ A, ord=pnp.inf),
             )
             subspace = (*shape_a, 2 ** len(wires))
 
-        A = qml.math.array(A) / normalization if normalization > 1 else A
+        # Clip the normalization to at least 1 (= normalize(A) if norm > 1 else A).
+        A = qml.math.array(A) / qml.math.maximum(normalization, qml.math.ones_like(normalization))
 
         if subspace[2] < (subspace[0] + subspace[1]):
             raise ValueError(


Test patch:
diff --git a/tests/ops/qubit/test_matrix_ops.py b/tests/ops/qubit/test_matrix_ops.py
index 1922f637a..5604dd012 100644
--- a/tests/ops/qubit/test_matrix_ops.py
+++ b/tests/ops/qubit/test_matrix_ops.py
@@ -795,6 +795,7 @@ class TestBlockEncode:
             (1, 1, {"norm": 1, "subspace": (1, 1, 2)}),
             ([1], 1, {"norm": 1, "subspace": (1, 1, 2)}),
             ([[1]], 1, {"norm": 1, "subspace": (1, 1, 2)}),
+            ([1, 0], [0, 1], {"norm": 1, "subspace": (1, 2, 4)}),
             (pnp.array(1), [1], {"norm": 1, "subspace": (1, 1, 2)}),
             (pnp.array([1]), 1, {"norm": 1, "subspace": (1, 1, 2)}),
             (pnp.array([[1]]), 1, {"norm": 1, "subspace": (1, 1, 2)}),
@@ -963,7 +964,7 @@ class TestBlockEncode:
         ("input_matrix", "wires", "output_matrix"),
         [
             (1, 0, [[1, 0], [0, -1]]),
-            (0.3, 0, [[0.3, 0.9539392], [0.9539392, -0.3]]),
+            ([0.3], 0, [[0.3, 0.9539392], [0.9539392, -0.3]]),
             (
                 [[0.1, 0.2], [0.3, 0.4]],
                 range(2),
@@ -983,6 +984,7 @@ class TestBlockEncode:
     )
     def test_blockencode_jax(self, input_matrix, wires, output_matrix):
         """Test that the BlockEncode operator matrix is correct for jax."""
+        import jax
         import jax.numpy as jnp
 
         input_matrix = jnp.array(input_matrix)
@@ -990,6 +992,14 @@ class TestBlockEncode:
         assert np.allclose(qml.matrix(op), output_matrix)
         assert qml.math.get_interface(qml.matrix(op)) == "jax"
 
+        # Test jitting behaviour as well.
+        @jax.jit
+        def f(A):
+            op = qml.BlockEncode(A, wires)
+            return qml.matrix(op)
+
+        assert np.allclose(f(input_matrix), output_matrix)
+
     @pytest.mark.parametrize("method", ["backprop"])
     @pytest.mark.parametrize(
         (