Batch run

doc/releases/changelog-dev.md

@@ -209,6 +209,9 @@
   Circuit fragments that are disconnected from the terminal measurements are now removed.
   [(#2254)](https://github.com/PennyLaneAI/pennylane/pull/2254)
 
+* A method to batch the non-trainable inputs for machine learning applications has been added. 
+[(#2069)](https://github.com/PennyLaneAI/pennylane/pull/2069)
+
 <h3>Improvements</h3>
 
 * The `qml.gradients` module has been streamlined and special-purpose functions
@@ -424,7 +427,7 @@ The Operator class has undergone a major refactor with the following changes:
 
 This release contains contributions from (in alphabetical order):
 
-Thomas Bromley, Anthony Hayes, Josh Izaac, Christina Lee,
+Jack Y. Araz, Thomas Bromley, Anthony Hayes, Josh Izaac, Christina Lee,
 Maria Fernanda Morris, Romain Moyard, Zeyue Niu, Maria Schuld, Jay Soni,
 Antal Száva, David Wierichs
 

pennylane/__init__.py

@@ -56,6 +56,7 @@
     adjoint,
     adjoint_metric_tensor,
     batch_params,
+    batch_input,
     batch_transform,
     draw,
     draw_mpl,

pennylane/qnn/keras.py

@@ -1,11 +1,11 @@
 # Copyright 2018-2021 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.
@@ -15,7 +15,9 @@
 API."""
 import inspect
 from collections.abc import Iterable
-from typing import Optional
+from typing import Optional, Union, Sequence, Text
+
+from pennylane.transforms.batch_input import batch_input
 
 try:
     import tensorflow as tf
@@ -52,6 +54,10 @@ class KerasLayer(Layer):
             arguments of the `add_weight()
             <https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer#add_weight>`__
             method and values being the corresponding specification.
+        batch_idx (Union[Sequence[int], int]): Argument location of the non-trainable inputs for
+            the circuit. This allows batch execution by creating executable circuits for each
+            input example with the same trainable weights. Default ``None``.
+            See :class:`~.pennylane.transforms.batch_input` for more details.
         **kwargs: additional keyword arguments passed to the Layer_ base class
 
     **Example**
@@ -196,8 +202,15 @@ def qnode(inputs, weights):
     """
 
     def __init__(
-        self, qnode, weight_shapes: dict, output_dim, weight_specs: Optional[dict] = None, **kwargs
+        self,
+        qnode,
+        weight_shapes: dict,
+        output_dim,
+        weight_specs: Optional[dict] = None,
+        batch_idx: Union[Sequence[int], int] = None,
+        **kwargs,
     ):
+        # pylint: disable=too-many-arguments
         if not CORRECT_TF_VERSION:
             raise ImportError(
                 "KerasLayer requires TensorFlow version 2 or above. The latest "
@@ -212,7 +225,12 @@ def __init__(
         }
 
         self._signature_validation(qnode, weight_shapes)
-        self.qnode = qnode
+
+        self.argnum = batch_idx
+        if batch_idx is None:
+            self.qnode = qnode
+        else:
+            self.qnode = batch_input(qnode, argnum=batch_idx)
 
         dtype = tf.float32 if tf.keras.backend.floatx() == tf.float32 else tf.float64
 
@@ -281,7 +299,7 @@ def call(self, inputs):
         Returns:
             tensor: output data
         """
-        if len(tf.shape(inputs)) > 1:
+        if len(tf.shape(inputs)) > 1 and self.argnum is None:
             # If the input size is not 1-dimensional, unstack the input along its first dimension,
             # recursively call the forward pass on each of the yielded tensors, and then stack the
             # outputs back into the correct shape
@@ -301,7 +319,10 @@ def _evaluate_qnode(self, x):
         Returns:
             tensor: output datapoint
         """
-        kwargs = {**{self.input_arg: x}, **{k: 1.0 * w for k, w in self.qnode_weights.items()}}
+        kwargs = {
+            **{self.input_arg: x},
+            **{k: 1.0 * w for k, w in self.qnode_weights.items()},
+        }
         return self.qnode(**kwargs)
 
     def compute_output_shape(self, input_shape):
@@ -330,3 +351,32 @@ def input_arg(self):
         `Layer <https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer>`__. Set to
         ``"inputs"``."""
         return self._input_arg
+
+    @staticmethod
+    def set_input_argument(input_name: Text = "inputs") -> None:
+        """
+        Set the name of the input argument.
+
+        Args:
+            input_name (str): Name of the input argument
+        """
+        KerasLayer._input_arg = input_name
+
+    def get_config(self) -> dict:
+        """
+        Get serialized layer configuration
+
+        Returns:
+            dict: layer configuration
+        """
+        config = super().get_config()
+
+        config.update(
+            {
+                "output_dim": self.output_dim,
+                "weight_specs": self.weight_specs,
+                "weight_shapes": self.weight_shapes,
+                "argnum": self.argnum,
+            }
+        )
+        return config

pennylane/transforms/__init__.py

@@ -31,6 +31,7 @@
 
     ~transforms.classical_jacobian
     ~batch_params
+    ~batch_input
     ~draw
     ~draw_mpl
     ~transforms.get_unitary_matrix
@@ -158,6 +159,7 @@
 from .op_transforms import op_transform
 from .adjoint import adjoint
 from .batch_params import batch_params
+from .batch_input import batch_input
 from .classical_jacobian import classical_jacobian
 from .compile import compile
 from .control import ControlledOperation, ctrl

pennylane/transforms/batch_input.py

@@ -0,0 +1,104 @@
+# Copyright 2018-2022 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.
+"""
+Batch transformation for multiple (non-trainable) input examples following issue #2037
+"""
+from typing import Union, Sequence, Callable, Tuple
+
+import pennylane as qml
+from pennylane import numpy as np
+from pennylane.transforms.batch_transform import batch_transform
+
+
+@batch_transform
+def batch_input(
+    tape: Union[qml.tape.JacobianTape, qml.QNode],
+    argnum: Union[Sequence[int], int] = 0,
+) -> Tuple[Sequence[qml.tape.JacobianTape], Callable]:
+    """
+    Transform a QNode to support an initial batch dimension for gate inputs.
+    In a classical ML application one needs to batch the non-trainable inputs of the network.
+    This function executes the same analogue for a quantum circuit:
+    separate circuit executions are created for each input, which are then executed
+    with the *same* trainable parameters.
+
+    The batch dimension is assumed to be the first rank of the
+    non trainable tensor object. For a rank 1 feature space, the shape needs to be ``(Nt, x)``
+    where ``x`` indicates the dimension of the features and ``Nt`` being the number of examples
+    within a batch.
+    Based on `arXiv:2202.10471 <https://arxiv.org/abs/2202.10471>__.
+
+    Args:
+        tape (.tape.JacobianTape or .QNode): Input quantum circuit to batch
+        argnum (Sequence[int] or int): One or more index value on all gate parameters
+            indicating the location of the non-trainable batched inputs within the input
+            argument sequence of the circuit. By default first argument is assumed to be
+            the only batched input.
+
+    Returns:
+        Sequence[Sequence[.tape.JacobianTape], Callable]: list of tapes arranged
+        according to unbatched inputs and a callable function to batch the results.
+
+    **Example**
+
+    .. code-block:: python
+
+            dev = qml.device("default.qubit", wires = 2, shots=None)
+
+            @batch_input(argnum=0)
+            @qml.qnode(dev, diff_method="parameter-shift", interface="tf")
+            def circuit(inputs, weights):
+                qml.AngleEmbedding(inputs, wires = range(2), rotation="Y")
+                qml.RY(weights[0], wires=0)
+                qml.RY(weights[1], wires=1)
+                return qml.expval(qml.PauliZ(1))
+
+        >>> x = np.random.uniform(0,1,(10,2))
+        >>> x.requires_grad = False
+        >>> w = np.random.uniform(0,1,2)
+        >>> circuit(x, w)
+        <tf.Tensor: shape=(10,), dtype=float64, numpy=
+        array([0.17926078, 0.7480163 , 0.47816999, 0.50381628, 0.349178  ,
+               0.17511444, 0.03769436, 0.19180259, 0.75867188, 0.55335748])>
+
+    .. seealso:: :func:`~.batch_params`
+    """
+    parameters = tape.get_parameters(trainable_only=False)
+
+    argnum = tuple(argnum) if isinstance(argnum, (list, tuple)) else (int(argnum),)
+
+    non_trainable, trainable = [], []
+    for idx, param in enumerate(parameters):
+        if idx in argnum:
+            non_trainable.append(param)
+        else:
+            trainable.append(param)
+
+    if len(np.unique([qml.math.shape(x)[0] for x in non_trainable])) != 1:
+        raise ValueError(
+            "Batch dimension for all gate arguments "
+            "specified by 'argnum' must be the same."
+        )
+
+    outputs = []
+    for inputs in zip(*non_trainable):
+        outputs += [list(inputs) + trainable]
+
+    # Construct new output tape with unstacked inputs
+    output_tapes = []
+    for params in outputs:
+        new_tape = tape.copy(copy_operations=True)
+        new_tape.set_parameters(params, trainable_only=False)
+        output_tapes.append(new_tape)
+
+    return output_tapes, lambda x: qml.math.squeeze(qml.math.stack(x))

tests/qnn/test_keras.py

@@ -615,3 +615,21 @@ def test_no_attribute():
     """Test that the qnn module raises an AttributeError if accessing an unavailable attribute"""
     with pytest.raises(AttributeError, match="module 'pennylane.qnn' has no attribute 'random'"):
         qml.qnn.random
+
+
+def test_batch_input():
+    """Test input batching in keras"""
+    dev = qml.device("default.qubit.tf", wires=4)
+
+    @qml.qnode(dev, diff_method="parameter-shift")
+    def circuit(x, weights):
+        qml.AngleEmbedding(x, wires=range(4), rotation="Y")
+        qml.RY(weights[0], wires=0)
+        qml.RY(weights[1], wires=1)
+        return qml.probs(op=qml.PauliZ(1))
+
+    KerasLayer.set_input_argument("x")
+    layer = KerasLayer(circuit, weight_shapes={"weights": (2,)}, output_dim=(2,), batch_idx=0)
+    conf = layer.get_config()
+    layer.build((None, 2))
+    assert layer(np.random.uniform(0, 1, (10, 4))).shape == (10, 2)

tests/transforms/test_batch_inputs.py

@@ -0,0 +1,71 @@
+# Copyright 2018-2021 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 the batch inputs transform.
+"""
+import functools
+import pytest
+
+import pennylane as qml
+from pennylane import numpy as np
+
+
+def test_simple_circuit():
+    """Test that batching works for a simple circuit"""
+    dev = qml.device("default.qubit", wires=2)
+
+    @qml.batch_input(argnum=0)
+    @qml.qnode(dev, diff_method="parameter-shift")
+    def circuit(inputs, weights):
+        qml.AngleEmbedding(inputs, wires=range(2), rotation="Y")
+        qml.RY(weights[0], wires=0)
+        qml.RY(weights[1], wires=1)
+        return qml.expval(qml.PauliZ(1))
+
+    batch_size = 5
+    inputs = np.random.uniform(0, np.pi, (batch_size, 2))
+    inputs.requires_grad = False
+    weights = np.random.uniform(-np.pi, np.pi, (2,))
+
+    res = circuit(inputs, weights)
+    assert res.shape == (batch_size,)
+
+
+def test_value_error():
+    """Test if the batch_input raises relevant errors correctly"""
+
+    dev = qml.device("default.qubit", wires=2)
+
+    @qml.batch_input(argnum=0)
+    @qml.qnode(dev, diff_method="parameter-shift")
+    def circuit(input1, input2, weights):
+        qml.AngleEmbedding(input1, wires=range(2), rotation="Y")
+        qml.RY(input2[0], wires=0)
+        qml.RY(weights[0], wires=0)
+        qml.RY(weights[1], wires=1)
+        return qml.expval(qml.PauliZ(1))
+
+    batch_size = 5
+    input1 = np.random.uniform(0, np.pi, (batch_size, 2))
+    input1.requires_grad = False
+    input2 = np.random.uniform(0, np.pi, (4, 1))
+    input2.requires_grad = False
+    weights = np.random.uniform(-np.pi, np.pi, (2,))
+
+    try:
+        res = circuit(input1, input2, weights)
+    except ValueError as err:
+        pass
+    except Exception as err:
+        raise Exception(err)