Repository: PennyLaneAI/pennylane
Pull request #5957 (2024-07-05 19:22:47)

Registers feature
**Context:**
Although pennylane allows the definition of wires, it doesn't easily support the act of creating whole registers (e.g. some ancillary register and measurement register). 

**Description of the Change:**
To make this process easier, this PR implements a new function called `qml.registers`, which takes in a dictionary as input, whose values are integers (number of qubits in the register). `qml.registers` will return a dictionary of Wires instances, where the key is the register name and the value is the Wires instance. `qml.registers` will also take, as input, nested registers of registers.

**Benefits:**
Implement registers much easier/faster on an abstract level. No need to write a for loop just to set up your Wires.

**Possible Drawbacks:**

**Related GitHub Issues:**


austingmhuang: [sc-67916]
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5957?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.66%. Comparing base [(`b7c25d4`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/b7c25d421adab5eb7a1515477a02fafccd41170b?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`1b04ec3`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/1b04ec333c51fce5c08826ce1c908ff3c1fa7c58?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    #5957      +/-   ##
==========================================
- Coverage   99.67%   99.66%   -0.01%     
==========================================
  Files         426      426              
  Lines       40871    40599     -272     
==========================================
- Hits        40738    40465     -273     
- Misses        133      134       +1     
```

</details>

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


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index a4c2b5d8d..089d8db61 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -15,8 +15,11 @@
   [(#5756)](https://github.com/PennyLaneAI/pennylane/pull/5756)
   [(#5987)](https://github.com/PennyLaneAI/pennylane/pull/5987)
 
-* The `split_to_single_terms` transform is added. This transform splits expectation values of sums
-  into multiple single-term measurements on a single tape, providing better support for simulators
+* A new function `qml.registers` has been added, enabling the creation of registers, which are implemented as a dictionary of `Wires` instances.
+  [(#5957)](https://github.com/PennyLaneAI/pennylane/pull/5957)
+
+* The `split_to_single_terms` transform is added. This transform splits expectation values of sums 
+  into multiple single-term measurements on a single tape, providing better support for simulators 
   that can handle non-commuting observables but don't natively support multi-term observables.
   [(#5884)](https://github.com/PennyLaneAI/pennylane/pull/5884)
 
@@ -235,12 +238,12 @@ Ahmed Darwish,
 Lillian M. A. Frederiksen,
 Pietropaolo Frisoni,
 Emiliano Godinez,
+Austin Huang,
 Renke Huang,
 Josh Izaac,
 Soran Jahangiri,
 Korbinian Kottmann,
 Christina Lee,
-Austin Huang,
 William Maxwell,
 Vincent Michaud-Rioux,
 Anurav Modak,
diff --git a/pennylane/__init__.py b/pennylane/__init__.py
index 8bc01c74f..3f0c484fc 100644
--- a/pennylane/__init__.py
+++ b/pennylane/__init__.py
@@ -60,7 +60,7 @@ from pennylane.circuit_graph import CircuitGraph
 from pennylane.configuration import Configuration
 from pennylane.drawer import draw, draw_mpl
 from pennylane.tracker import Tracker
-
+from pennylane.registers import registers
 from pennylane.io import *
 from pennylane.measurements import (
     counts,
diff --git a/pennylane/registers.py b/pennylane/registers.py
new file mode 100644
index 000000000..7b710491f
--- /dev/null
+++ b/pennylane/registers.py
@@ -0,0 +1,124 @@
+# Copyright 2018-2024 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+This module contains the :func:`registers` function.
+"""
+
+from .wires import Wires
+
+
+def registers(register_dict):
+    """
+    Returns a dictionary that maps register names to :class:`~.Wires`.
+
+    This function helps to group qubits and abstract away the finer details of running quantum
+    algorithms. Register names and their total number of wires are typically known in advance,
+    but managing the specific wire range for each register can be a challenge. The ``qml.registers()``
+    function creates a dictionary that maps register name to :class:`~.Wires` object. Moreover,
+    it allows one to input a nested structure where registers contain sub-registers, as illustrated
+    in the examples below.
+
+    Args:
+        register_dict (dict): a dictionary where keys are register names and values are either
+            positive integers indicating the number of qubits or nested dictionaries of more registers.
+
+    Returns:
+        dict (Wires): dictionary where the keys are the names (str) of the registers, and the
+        values are :class:`~.Wires` objects.
+
+    **Example**
+
+    Given flat input dictionary:
+
+    >>> qml.registers({"alice": 2, "bob": 3})
+    {'alice': Wires([0, 1]), 'bob': Wires([2, 3, 4])}
+
+    Given nested input dictionary:
+
+    >>> wire_registers = qml.registers({"people": {"alice": 2, "bob": 1}})
+    >>> wire_dict
+    {'alice': Wires([0, 1]), 'bob': Wires([2]), 'people': Wires([0, 1, 2])}
+    >>> wire_dict['bob']
+    Wires([2])
+    >>> wire_dict['alice'][1]
+    1
+
+    A simple example showcasing how to implement the `SWAP <https://en.wikipedia.org/wiki/Swap_test>`_ test:
+
+    .. code-block::
+
+        dev = qml.device("default.qubit")
+        reg =  registers({"aux": 1, "phi": 5, "psi": 5})
+
+        @qml.qnode(dev)
+        def circuit():
+            for state in ["phi", "psi"]:
+                 qml.BasisState([1, 1, 0, 0, 0], reg[state])
+
+            qml.Hadamard(reg["aux"])
+            for i in range(len(reg["phi"])):
+                qml.CSWAP(reg["aux"] + reg["phi"][i] + reg["psi"][i])
+            qml.Hadamard(reg["aux"])
+
+            return qml.expval(qml.Z(wires=reg["aux"]))
+
+        >>> circuit()
+        0.9999999999999996
+    """
+
+    def _registers(register_dict, _start_wire_index=0):
+        """Recursively builds a dictionary of Wires objects from a nested dictionary of register
+        names and sizes.
+
+        Args:
+            register_dict (dict): a dictionary where keys are register names and values are either
+                positive integers indicating the number of qubits or nested dictionaries of more registers
+            _start_wire_index (int): the starting index for the wire labels.
+
+        Returns:
+            dict (Wires): dictionary where the keys are the names (str) of the registers, and the
+            values are either :class:`~.Wires` objects or other registers
+        """
+
+        all_reg = {}
+        for register_name, register_wires in register_dict.items():
+            if isinstance(register_wires, dict):
+                if len(register_wires) == 0:
+                    raise ValueError(f"Got an empty dictionary '{register_wires}'")
+                inner_register_dict = _registers(register_wires, _start_wire_index)
+                wire_vals = [w for reg in inner_register_dict.values() for w in reg.tolist()]
+                all_reg.update(inner_register_dict)
+                all_reg[register_name] = Wires(
+                    range(
+                        _start_wire_index,
+                        all_reg[next(reversed(inner_register_dict))].labels[-1] + 1,
+                    )
+                )
+                _start_wire_index = wire_vals[-1] + 1
+            elif isinstance(register_wires, int):
+                if register_wires < 1:
+                    raise ValueError(
+                        f"Expected '{register_wires}' to be greater than 0. Please ensure that "
+                        "the number of wires for the register is a positive integer"
+                    )
+                all_reg[register_name] = Wires(
+                    range(_start_wire_index, register_wires + _start_wire_index)
+                )
+                _start_wire_index += register_wires
+            else:  # Not a dict nor an int
+                raise ValueError(f"Expected '{register_wires}' to be either a dict or an int. ")
+
+        return all_reg
+
+    return _registers(register_dict)


Test patch:
diff --git a/tests/test_registers.py b/tests/test_registers.py
new file mode 100644
index 000000000..9a6593c7f
--- /dev/null
+++ b/tests/test_registers.py
@@ -0,0 +1,106 @@
+# Copyright 2018-2024 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Unit tests for pennylane.registers.
+"""
+
+import pytest
+
+import pennylane as qml
+from pennylane.wires import Wires
+
+
+class TestRegisters:
+    @pytest.mark.parametrize(
+        "wire_dict, expected_register",
+        [
+            (
+                {"alice": 3, "bob": {"nest1": 3, "nest2": 3}, "cleo": 3},
+                {
+                    "alice": Wires([0, 1, 2]),
+                    "nest1": Wires([3, 4, 5]),
+                    "nest2": Wires([6, 7, 8]),
+                    "bob": Wires([3, 4, 5, 6, 7, 8]),
+                    "cleo": Wires([9, 10, 11]),
+                },
+            ),
+            (
+                {"alice": 3, "bob": {"nest1": 3}, "cleo": 3},
+                {
+                    "alice": Wires([0, 1, 2]),
+                    "nest1": Wires([3, 4, 5]),
+                    "bob": Wires([3, 4, 5]),
+                    "cleo": Wires([6, 7, 8]),
+                },
+            ),
+            (
+                {"alice": 3, "bob": {"nest1": {"nest2": {"nest3": 1}}}, "cleo": 3},
+                {
+                    "alice": Wires([0, 1, 2]),
+                    "nest3": Wires([3]),
+                    "nest2": Wires([3]),
+                    "nest1": Wires([3]),
+                    "bob": Wires([3]),
+                    "cleo": Wires([4, 5, 6]),
+                },
+            ),
+            (
+                {
+                    "alice": 3,
+                    "bob": {"nest1": {"nest2": {"nest3": 1}, "nest2a": 2}, "nest1a": 3},
+                    "cleo": 3,
+                },
+                {
+                    "alice": Wires([0, 1, 2]),
+                    "nest3": Wires([3]),
+                    "nest2": Wires([3]),
+                    "nest2a": Wires([4, 5]),
+                    "nest1": Wires([3, 4, 5]),
+                    "nest1a": Wires([6, 7, 8]),
+                    "bob": Wires([3, 4, 5, 6, 7, 8]),
+                    "cleo": Wires([9, 10, 11]),
+                },
+            ),
+        ],
+    )
+    def test_build_registers(self, wire_dict, expected_register):
+        """Test that the registers function returns the correct dictionary of Wires instances.
+        The expected result is a dictionary that has elements ordered first by appearance from left
+        to right, then by nestedness (also known as DFS traversal order)."""
+
+        wire_dict = qml.registers(wire_dict)
+
+        assert wire_dict == expected_register
+
+    @pytest.mark.parametrize(
+        "wire_dict, expected_error_msg",
+        [
+            (
+                {"alice": 3, "bob": {"nest1": 3, "nest2": {}}, "cleo": 3},
+                "Got an empty dictionary",
+            ),
+            (
+                {"alice": 3, "bob": {"nest1": 0}, "cleo": 3},
+                "Expected '0' to be greater than 0. Please ensure that the number of wires for the register is a positive integer",
+            ),
+            (
+                {"alice": 3, "bob": {"nest1": {"nest2": {"nest3": {1}}}}, "cleo": 3},
+                r"Expected '\{1\}' to be either a dict or an int",
+            ),
+        ],
+    )
+    def test_errors_for_registers(self, wire_dict, expected_error_msg):
+        """Test that the registers function raises the right error for given Wires dictionaries"""
+        with pytest.raises(ValueError, match=expected_error_msg):
+            qml.registers(wire_dict)