Add Dataset Attribute type for Pytrees

doc/releases/changelog-dev.md

@@ -177,6 +177,10 @@
 * `QuantumScript` properties are only calculated when needed, instead of on initialization. This decreases the classical overhead by >20%.
   `par_info`, `obs_sharing_wires`, and `obs_sharing_wires_id` are now public attributes.
   [(#5696)](https://github.com/PennyLaneAI/pennylane/pull/5696)
+
+* The `qml.data` module now supports PyTree types as dataset attributes
+  [(#5732)](https://github.com/PennyLaneAI/pennylane/pull/5732)
+
 
 * `qml.ops.Conditional` now inherits from `qml.ops.SymbolicOp`, thus it inherits several useful common functionalities. Other properties such as adjoint and diagonalizing gates have been added using the `base` properties.
   [(##5772)](https://github.com/PennyLaneAI/pennylane/pull/5772)
@@ -342,6 +346,7 @@ This release contains contributions from (in alphabetical order):
 Guillermo Alonso-Linaje,
 Lillian M. A. Frederiksen,
 Gabriel Bottrill,
+Jack Brown,
 Astral Cai,
 Ahmed Darwish,
 Isaac De Vlugt,

pennylane/data/__init__.py

@@ -211,6 +211,7 @@ class QuantumOscillator(qml.data.Dataset, data_name="quantum_oscillator", identi
     DatasetSparseArray,
     DatasetString,
     DatasetTuple,
+    DatasetPyTree,
 )
 from .base import DatasetNotWriteableError
 from .base.attribute import AttributeInfo, DatasetAttribute, attribute
@@ -225,6 +226,7 @@ class QuantumOscillator(qml.data.Dataset, data_name="quantum_oscillator", identi
     "DatasetAttribute",
     "DatasetNotWriteableError",
     "DatasetArray",
+    "DatasetPyTree",
     "DatasetScalar",
     "DatasetString",
     "DatasetList",

pennylane/data/attributes/__init__.py

@@ -24,6 +24,7 @@
 from .sparse_array import DatasetSparseArray
 from .string import DatasetString
 from .tuple import DatasetTuple
+from .pytree import DatasetPyTree
 
 __all__ = (
     "DatasetArray",
@@ -32,6 +33,7 @@
     "DatasetDict",
     "DatasetList",
     "DatasetOperator",
+    "DatasetPyTree",
     "DatasetSparseArray",
     "DatasetMolecule",
     "DatasetNone",

pennylane/data/attributes/operator/operator.py

@@ -52,7 +52,8 @@ class DatasetOperator(Generic[Op], DatasetAttribute[HDF5Group, Op, Op]):
 
     @classmethod
     @lru_cache(1)
-    def consumes_types(cls) -> FrozenSet[Type[Operator]]:
+    def supported_ops(cls) -> FrozenSet[Type[Operator]]:
+        """Set of supported operators."""
         return frozenset(
             (
                 # pennylane/operation/Tensor
@@ -214,7 +215,7 @@ def _ops_to_hdf5(
             op_key = f"op_{i}"
             if isinstance(op, (qml.ops.Prod, qml.ops.SProd, qml.ops.Sum)):
                 op = op.simplify()
-            if type(op) not in self.consumes_types():
+            if type(op) not in self.supported_ops():
                 raise TypeError(
                     f"Serialization of operator type '{type(op).__name__}' is not supported."
                 )
@@ -254,6 +255,7 @@ def _hdf5_to_ops(self, bind: HDF5Group) -> List[Operator]:
         wires_bind = bind["op_wire_labels"]
         op_class_names = [] if names_bind.shape == (0,) else names_bind.asstr()
         op_wire_labels = [] if wires_bind.shape == (0,) else wires_bind.asstr()
+
         with qml.QueuingManager.stop_recording():
             for i, op_class_name in enumerate(op_class_names):
                 op_key = f"op_{i}"
@@ -293,4 +295,4 @@ def _hdf5_to_ops(self, bind: HDF5Group) -> List[Operator]:
     @lru_cache(1)
     def _supported_ops_dict(cls) -> Dict[str, Type[Operator]]:
         """Returns a dict mapping ``Operator`` subclass names to the class."""
-        return {op.__name__: op for op in cls.consumes_types()}
+        return {op.__name__: op for op in cls.supported_ops()}

pennylane/data/attributes/pytree.py

@@ -0,0 +1,57 @@
+# 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.
+"""Contains DatasetAttribute definition for PyTree types."""
+
+
+from typing import TypeVar
+
+import numpy as np
+
+from pennylane.data.attributes import DatasetArray, DatasetList
+from pennylane.data.base.attribute import DatasetAttribute
+from pennylane.data.base.hdf5 import HDF5Group
+from pennylane.data.base.mapper import AttributeTypeMapper
+from pennylane.pytrees import flatten, serialization, unflatten
+
+T = TypeVar("T")
+
+
+class DatasetPyTree(DatasetAttribute[HDF5Group, T, T]):
+    """Attribute type for an object that can be converted to
+    a Pytree. This is the default serialization method for
+    all Pennylane Pytrees, including subclasses of ``Operator``.
+    """
+
+    type_id = "pytree"
+
+    def hdf5_to_value(self, bind: HDF5Group) -> T:
+        return unflatten(
+            AttributeTypeMapper(bind)["leaves"].get_value(),
+            serialization.pytree_structure_load(bind["treedef"][()].tobytes()),
+        )
+
+    def value_to_hdf5(self, bind_parent: HDF5Group, key: str, value: T) -> HDF5Group:
+        bind = bind_parent.create_group(key)
+        leaves, treedef = flatten(value)
+
+        bind["treedef"] = np.void(serialization.pytree_structure_dump(treedef, decode=False))
+
+        try:
+            # Attempt to store leaves as an array, which will be more efficient
+            # but will fail if the leaves are not homogenous
+            DatasetArray(leaves, parent_and_key=(bind, "leaves"))
+        except (ValueError, TypeError):
+            DatasetList(leaves, parent_and_key=(bind, "leaves"))
+
+        return bind

pennylane/data/base/attribute.py

@@ -38,6 +38,7 @@
 from pennylane.data.base import hdf5
 from pennylane.data.base.hdf5 import HDF5, HDF5Any, HDF5Group
 from pennylane.data.base.typing_util import UNSET, get_type, get_type_str
+from pennylane.pytrees import is_pytree
 
 T = TypeVar("T")
 
@@ -492,5 +493,7 @@ def match_obj_type(
         ret = DatasetAttribute.registry["list"]
     elif issubclass(type_, Mapping):
         ret = DatasetAttribute.registry["dict"]
+    elif is_pytree(type_):
+        ret = DatasetAttribute.registry["pytree"]
 
     return ret

pennylane/measurements/shots.py

@@ -11,8 +11,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """This module contains the Shots class to hold shot-related information."""
+from collections.abc import Sequence
+
 # pylint:disable=inconsistent-return-statements
-from typing import NamedTuple, Sequence, Tuple
+from typing import NamedTuple
 
 
 class ShotCopies(NamedTuple):
@@ -39,7 +41,7 @@ def valid_int(s):
 
 def valid_tuple(s):
     """Returns True if s is a tuple of the form (shots, copies)."""
-    return isinstance(s, tuple) and len(s) == 2 and valid_int(s[0]) and valid_int(s[1])
+    return isinstance(s, Sequence) and len(s) == 2 and valid_int(s[0]) and valid_int(s[1])
 
 
 class Shots:
@@ -136,7 +138,7 @@ class Shots:
     total_shots: int = None
     """The total number of shots to be executed."""
 
-    shot_vector: Tuple[ShotCopies] = None
+    shot_vector: tuple[ShotCopies] = None
     """The tuple of :class:`~ShotCopies` to be executed. Each element is of the form ``(shots, copies)``."""
 
     _SHOT_ERROR = ValueError(
@@ -167,7 +169,7 @@ def __init__(self, shots=None):
         elif isinstance(shots, Sequence):
             if not all(valid_int(s) or valid_tuple(s) for s in shots):
                 raise self._SHOT_ERROR
-            self.__all_tuple_init__([s if isinstance(s, tuple) else (s, 1) for s in shots])
+            self.__all_tuple_init__([s if isinstance(s, Sequence) else (s, 1) for s in shots])
         elif isinstance(shots, self.__class__):
             return  # self already _is_ shots as defined by __new__
         else:
@@ -211,7 +213,7 @@ def __iter__(self):
             for _ in range(shot_copy.copies):
                 yield shot_copy.shots
 
-    def __all_tuple_init__(self, shots: Sequence[Tuple]):
+    def __all_tuple_init__(self, shots: Sequence[tuple]):
         res = []
         total_shots = 0
         current_shots, current_copies = shots[0]

pennylane/pytrees/__init__.py

@@ -0,0 +1,27 @@
+# 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.
+"""
+An internal module for working with pytrees.
+"""
+
+from .pytrees import PyTreeStructure, flatten, is_pytree, leaf, register_pytree, unflatten
+
+__all__ = [
+    "PyTreeStructure",
+    "flatten",
+    "is_pytree",
+    "leaf",
+    "register_pytree",
+    "unflatten",
+]