QROM subroutine func decomp

doc/releases/changelog-dev.md

@@ -53,6 +53,7 @@
   [(#7779)](https://github.com/PennyLaneAI/pennylane/pull/7779)
   [(#7908)](https://github.com/PennyLaneAI/pennylane/pull/7908)
   [(#7385)](https://github.com/PennyLaneAI/pennylane/pull/7385)
+  [(#7941)](https://github.com/PennyLaneAI/pennylane/pull/7941)
 
   The included templates are:
 
@@ -86,6 +87,8 @@
 
   * :class:`~.TrotterProduct`
 
+  * :class:`~.QROM`
+
 * A new function called :func:`~.math.choi_matrix` is available, which computes the [Choi matrix](https://en.wikipedia.org/wiki/Choi%E2%80%93Jamio%C5%82kowski_isomorphism) of a quantum channel.
   This is a useful tool in quantum information science and to check circuit identities involving non-unitary operations.
   [(#7951)](https://github.com/PennyLaneAI/pennylane/pull/7951)

pennylane/templates/subroutines/qrom.py

@@ -16,10 +16,18 @@
 """
 
 import math
+from collections import Counter
+from functools import reduce
 
 import numpy as np
 
 from pennylane import ops as qml_ops
+from pennylane.decomposition import (
+    add_decomps,
+    controlled_resource_rep,
+    register_resources,
+    resource_rep,
+)
 from pennylane.operation import Operation
 from pennylane.queuing import QueuingManager, apply
 from pennylane.templates.embeddings import BasisEmbedding
@@ -109,6 +117,14 @@
 
     """
 
+    resource_keys = {
+        "num_bitstrings",
+        "num_control_wires",
+        "num_target_wires",
+        "num_work_wires",
+        "clean",
+    }
+
     def __init__(
         self,
         bitstrings,
@@ -157,6 +173,16 @@
         metadata = tuple((key, value) for key, value in self.hyperparameters.items())
         return tuple(), metadata
 
+    @property
+    def resource_params(self) -> dict:
+        return {
+            "num_bitstrings": len(self.hyperparameters["bitstrings"]),
+            "num_control_wires": len(self.hyperparameters["control_wires"]),
+            "num_target_wires": len(self.hyperparameters["target_wires"]),
+            "num_work_wires": len(self.hyperparameters["work_wires"]),
+            "clean": self.hyperparameters["clean"],
+        }
+
     @classmethod
     def _unflatten(cls, data, metadata):
         hyperparams_dict = dict(metadata)
@@ -310,3 +336,190 @@
     def clean(self):
         """Boolean to select the version of QROM."""
         return self.hyperparameters["clean"]
+
+
+def _qrom_decomposition_resources(
+    num_bitstrings, num_control_wires, num_target_wires, num_work_wires, clean
+):
+    if num_control_wires == 0:
+        return {resource_rep(BasisEmbedding, num_wires=num_target_wires): num_bitstrings}
+
+    num_swap_wires = num_target_wires + num_work_wires
+
+    # number of operators we store per column (power of 2)
+    depth = num_swap_wires // num_target_wires
+    depth = int(2 ** np.floor(np.log2(depth)))
+    depth = min(depth, num_bitstrings)
+
+    ops = [resource_rep(BasisEmbedding, num_wires=num_target_wires) for _ in range(num_bitstrings)]
+    ops_identity = ops + [qml_ops.I] * int(2**num_control_wires - num_bitstrings)
+
+    n_columns = (
+        num_bitstrings // depth if num_bitstrings % depth == 0 else num_bitstrings // depth + 1
+    )
+
+    # New ops block
+    new_ops = Counter()
+    for i in range(n_columns):
+        column_ops = Counter()
+        for j in range(depth):
+            column_ops[ops_identity[i * depth + j]] += 1
+        if len(column_ops) == 1 and list(column_ops.values())[0] == 1:
+            new_ops[list(column_ops.keys())[0]] += 1
+        else:
+            new_ops[resource_rep(qml_ops.op_math.Prod, resources=dict(column_ops))] += 1
+
+    # Select block
+    num_control_select_wires = int(math.ceil(math.log2(2**num_control_wires / depth)))
+
+    new_ops_reps = reduce(
+        lambda acc, lst: acc + lst, [[key for _ in range(val)] for key, val in new_ops.items()]
+    )
+
+    if num_control_select_wires > 0:
+        select_ops = {
+            resource_rep(
+                Select, num_control_wires=num_control_select_wires, op_reps=tuple(new_ops_reps)
+            ): 1
+        }
+    else:
+        select_ops = new_ops
+
+    # Swap block
+    num_control_swap_wires = num_control_wires - num_control_select_wires
+    swap_resources = Counter()
+    for ind in range(num_control_swap_wires):
+        for j in range(2**ind):
+            num_swaps = min(
+                (j + 1) * num_target_wires - (j) * num_target_wires,
+                (j + 2 ** (ind + 1)) * num_target_wires - (j + 2**ind) * num_target_wires,
+            )
+            if num_swaps > 1:
+                swaps = {qml_ops.SWAP: num_swaps}
+                swap_resources[
+                    controlled_resource_rep(
+                        base_class=qml_ops.op_math.Prod,
+                        base_params={"resources": swaps},
+                        num_control_wires=1,
+                    )
+                ] += 1
+            else:
+                swap_resources[
+                    controlled_resource_rep(
+                        base_class=qml_ops.SWAP,
+                        base_params={},
+                        num_control_wires=1,
+                    )
+                ] += 1
+
+    if not clean or depth == 1:
+        resources = swap_resources
+        resources.update(select_ops)
+        return resources
+
+    resources = {}
+
+    hadamard_ops = {qml_ops.Hadamard: num_target_wires}
+
+    for key, val in swap_resources.items():
+        swap_resources[key] = val * 2
+
+    resources.update(hadamard_ops)
+    resources.update(swap_resources)
+    resources.update(select_ops)
+
+    for key, val in resources.items():
+        resources[key] = val * 2
+
+    return resources
+
+
+@register_resources(_qrom_decomposition_resources)
+def _qrom_decomposition(
+    wires, bitstrings, control_wires, target_wires, work_wires, clean
+):  # pylint: disable=unused-argument, too-many-arguments
+    if len(control_wires) == 0:
+        for bits in bitstrings:
+            BasisEmbedding(int(bits, 2), wires=target_wires)
+
+    swap_wires = target_wires + work_wires
+
+    # number of operators we store per column (power of 2)
+    depth = len(swap_wires) // len(target_wires)
+    depth = int(2 ** np.floor(np.log2(depth)))
+    depth = min(depth, len(bitstrings))
+
+    def _new_ops(depth_new, target_wires_new):
+
+        with QueuingManager.stop_recording():
+            ops_new = [BasisEmbedding(int(bits, 2), wires=target_wires) for bits in bitstrings]
+            ops_identity_new = ops_new + [qml_ops.I(target_wires)] * int(
+                2 ** len(control_wires) - len(ops_new)
+            )
+
+        n_columns = (
+            len(bitstrings) // depth_new
+            if len(bitstrings) % depth_new == 0
+            else len(bitstrings) // depth_new + 1
+        )
+        new_ops = []
+        for i in range(n_columns):
+            column_ops = []
+            for j in range(depth_new):
+                dic_map = {
+                    ops_identity_new[i * depth_new + j].wires[l]: swap_wires[
+                        j * len(target_wires_new) + l
+                    ]
+                    for l in range(len(target_wires))
+                }
+                column_ops.append(ops_identity_new[i * depth_new + j].map_wires(dic_map))
+            new_ops.append(qml_ops.prod(*column_ops))
+        return new_ops
+
+    def _select_ops(control_wires_select, depth_select, target_wires_select):
+        n_control_select_wires = int(
+            math.ceil(math.log2(2 ** len(control_wires_select) / depth_select))
+        )
+        control_select_wires = control_wires_select[:n_control_select_wires]
+
+        if control_select_wires:
+            Select(
+                _new_ops(depth_select, target_wires_select),
+                control=control_select_wires,
+            )
+        else:
+            _new_ops(depth_select, target_wires_select)
+
+    def _swap_ops(control_wires_swap, depth_swap, swap_wires_swap, target_wires_swap):
+        n_control_select_wires = int(
+            math.ceil(math.log2(2 ** len(control_wires_swap) / depth_swap))
+        )
+        control_swap_wires = control_wires_swap[n_control_select_wires:]
+        for ind in range(len(control_swap_wires)):
+            for j in range(2**ind):
+                new_op = qml_ops.prod(_multi_swap)(
+                    swap_wires_swap[
+                        (j) * len(target_wires_swap) : (j + 1) * len(target_wires_swap)
+                    ],
+                    swap_wires_swap[
+                        (j + 2**ind)
+                        * len(target_wires_swap) : (j + 2 ** (ind + 1))
+                        * len(target_wires_swap)
+                    ],
+                )
+                qml_ops.ctrl(new_op, control=control_swap_wires[-ind - 1])
+
+    if not clean or depth == 1:
+        _select_ops(control_wires, depth, target_wires)
+        _swap_ops(control_wires, depth, swap_wires, target_wires)
+
+    else:
+        for _ in range(2):
+            for w in target_wires:
+                qml_ops.Hadamard(wires=w)
+                _swap_ops(control_wires, depth, swap_wires, target_wires)
+                _select_ops(control_wires, depth, target_wires)
+                _swap_ops(control_wires, depth, swap_wires, target_wires)
+
+
+add_decomps(QROM, _qrom_decomposition)

tests/templates/test_subroutines/test_qrom.py

@@ -19,6 +19,7 @@
 
 import pennylane as qml
 from pennylane import numpy as np
+from pennylane.ops.functions.assert_valid import _test_decomposition_rule
 
 
 def test_assert_valid_qrom():
@@ -192,6 +193,28 @@ def test_decomposition(self):
         for op1, op2 in zip(qrom_decomposition, expected_gates):
             qml.assert_equal(op1, op2)
 
+    @pytest.mark.parametrize(
+        ("bitstrings", "control_wires", "target_wires", "work_wires", "clean"),
+        [
+            (["1", "0", "0", "1"], [0, 1], [2], [3], True),
+            (["10", "00", "00", "01", "01", "00", "00", "01"], [0, 1, 2], [3, 4], [5], False),
+            (["01", "00", "00", "10", "10", "00", "00", "01"], [0, 1, 2], [3, 4], [5], True),
+            (["1", "0", "0", "1"], [0, 1], [2], [], False),
+            (["1", "0", "0", "1"], [0, 1], [2], [3, 4], False),
+        ],  # pylint: disable=too-many-arguments
+    )
+    def test_decomposition_new(self, bitstrings, control_wires, target_wires, work_wires, clean):
+        """Tests the decomposition rule implemented with the new system."""
+        op = qml.QROM(
+            bitstrings,
+            control_wires=control_wires,
+            target_wires=target_wires,
+            work_wires=work_wires,
+            clean=clean,
+        )
+        for rule in qml.list_decomps(qml.QROM):
+            _test_decomposition_rule(op, rule)
+
     def test_zero_control_wires(self):
         """Test that the edge case of zero control wires works"""