Repository: PennyLaneAI/pennylane
Pull request #2069 (2021-12-27 15:34:16)

Batch run
**Context:**
In classical ML applications, training examples are run in batches and then the mean or sum of an objective function upon this batch is calculated to apply gradient descent. This function allows the collective execution of a batched input tape with the same weight used for the entire batch.

**Description of the Change:**
Additional transformation function.

**Benefits:**
A batch of training or validation examples will run at the same time. This will also allow collective job submissions to IBM-Q instead of one circuit tape at a time.

**Possible Drawbacks:**
Code assumes that the arguments are ordered with respect to the batched inputs first and then non-batched inputs. If the user requests any other order this will lead to wrong results.

**Related GitHub Issues:**
Improvement request mentioned in issue #2037 

josh146: Hi @jackaraz, thanks for making this PR! I just wanted to check in to see if there was anything you had any questions regarding. If not, would this be ready for a code review?
jackaraz: Hi @josh146 not at the moment, I guess my code does not meet all the requirements that you have about code formatting. I can modify it with black if you like. Also after seeing @antalszava 's proposal in [here](https://github.com/PennyLaneAI/pennylane/issues/2068#issuecomment-1007199481) I was thinking that might be possible to update the TensorFlow layer a bit further. However, this will require creating it as a Model rather than a layer which might create usability issues for a regular user.
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2069?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#2069](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2069?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (9331ec7) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/208095a4c9cb78e357cb9e1cae55939eeeeaeca1?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (208095a) will **decrease** coverage by `2.10%`.
> The diff coverage is `86.66%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2069/graphs/tree.svg?width=650&height=150&src=pr&token=9S3YzlUgGk&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2069?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@            Coverage Diff             @@
##           master    #2069      +/-   ##
==========================================
- Coverage   99.18%   97.08%   -2.11%     
==========================================
  Files         226      227       +1     
  Lines       17369    17406      +37     
==========================================
- Hits        17228    16899     -329     
- Misses        141      507     +366     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2069?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2069/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL19faW5pdF9fLnB5) | `100.00% <ø> (ø)` | |
| [pennylane/qnn/keras.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2069/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3Fubi9rZXJhcy5weQ==) | `93.67% <73.68%> (-6.33%)` | :arrow_down: |
| [pennylane/transforms/batch\_input.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2069/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvYmF0Y2hfaW5wdXQucHk=) | `96.00% <96.00%> (ø)` | |
| [pennylane/transforms/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2069/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvX19pbml0X18ucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/interfaces/batch/jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2069/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvamF4LnB5) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pennylane/interfaces/jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2069/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvamF4LnB5) | `0.00% <0.00%> (-97.44%)` | :arrow_down: |
| [pennylane/devices/default\_qubit\_jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2069/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9xdWJpdF9qYXgucHk=) | `0.00% <0.00%> (-96.25%)` | :arrow_down: |
| [pennylane/qnn/cost.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2069/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3Fubi9jb3N0LnB5) | `42.85% <0.00%> (-57.15%)` | :arrow_down: |
| [pennylane/transforms/classical\_jacobian.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2069/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvY2xhc3NpY2FsX2phY29iaWFuLnB5) | `80.00% <0.00%> (-20.00%)` | :arrow_down: |
| [pennylane/math/utils.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2069/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvdXRpbHMucHk=) | `81.81% <0.00%> (-18.19%)` | :arrow_down: |
| ... and [16 more](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2069/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2069?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2069?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [208095a...9331ec7](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2069?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

josh146: > I can modify it with black if you like.

Yes, that would be great 🙂 

> Also after seeing @antalszava 's proposal in here I was thinking that might be possible to update the TensorFlow layer a bit further. However, this will require creating it as a Model rather than a layer which might create usability issues for a regular user.

Feel free to discuss any thoughts you have regarding this in more detail, either here on in the issue! As you are currently working on a model that requires this feature, your feedback is invaluable.
jackaraz: Hi @josh146 

> > I can modify it with black if you like.
> 
> Yes, that would be great 🙂
This is done. But I still see some linting errors; I believe it doesn't like how I present `argnum` in the Keras layer. 

> > Also after seeing @antalszava 's proposal in here I was thinking that might be possible to update the TensorFlow layer a bit further. However, this will require creating it as a Model rather than a layer which might create usability issues for a regular user.
> 
> Feel free to discuss any thoughts you have regarding this in more detail, either here on in the issue! As you are currently working on a model that requires this feature, your feedback is invaluable.

I do have a few more implementations based on TF using pennylane, such as quantum natural gradient (as far as I know, this does not exist with TF backend at the moment, my implementation is based on pennylane's original one) and for parallelization on GPU/CPU. However, my implementation is not very generic; I'm basically writing a Keras model and a custom training within it. But it requires different models for different implementations i.e. one model for pure Quantum circuit-based networks, one for hybrid where the classical portion can train with traditional SGD and the quantum portion can train with QNGD within the same training sequence etc. However, as I said, these are case dependent implementations, we will release the paper soon, and I can show you the implementations, but I'm not sure if it would be relevant for pennylane.
josh146: Thanks @jackaraz!

> I do have a few more implementations based on TF using pennylane, such as quantum natural gradient (as far as I know, this does not exist with TF backend at the moment, my implementation is based on pennylane's original one)

This actually would be very interesting to see, as we have long been wanting to extend the QNG optimizer to support other interfaces. However, we still have various design questions, so your implementation --- even if not eventually merged into the codebase --- could still be helpful in resolving these questions!
vbelis: Hello everyone. I was preparing to open a new issue regarding the batching of input data. However, it seems that @jackaraz and @josh146 have already discussed this in the [forum](https://discuss.pennylane.ai/t/tensorflow-pennylane-ibm-quantum/1559), before starting this PR, and would like to ensure whether we have the same goal in mind.

The implementation of `tranform.batch_params`([source](https://github.com/PennyLaneAI/pennylane/blob/master/pennylane/transforms/batch_params.py)) adds a batch dimension to each argument that is trainable by default, or to all `QNode` arguments if `all_operations=True` is enabled. This is, in a sense, the "opposite" of what it is typically required in a neural network training workflow, where, it is the inputs (encoded data features) that have a batch dimension whereas the weights don't. Hence, the weight tensors should be kept constant throughout the forward pass whereas the input `x` should have dimensions/shape: `(n_batch, n_features)`. 

My understanding is that the goal of this PR is to implement the above functionality. I am wondering whether your implementation @jackaraz is `TensorFlow` specific. Have you tested whether it can be used without an `interface`, i.e. vanilla `pennylane`, or with `PyTorch` (`qnn.TorchLayer`)?


jackaraz: > Hello everyone. I was preparing to open a new issue regarding the batching of input data. However, it seems that @jackaraz and @josh146 have already discussed this in the [forum](https://discuss.pennylane.ai/t/tensorflow-pennylane-ibm-quantum/1559), before starting this PR, and would like to ensure whether we have the same goal in mind.
> 
> The implementation of `tranform.batch_params`([source](https://github.com/PennyLaneAI/pennylane/blob/master/pennylane/transforms/batch_params.py)) adds a batch dimension to each argument that is trainable by default, or to all `QNode` arguments if `all_operations=True` is enabled. This is, in a sense, the "opposite" of what it is typically required in a neural network training workflow, where, it is the inputs (encoded data features) that have a batch dimension whereas the weights don't. Hence, the weight tensors should be kept constant throughout the forward pass whereas the input `x` should have dimensions/shape: `(n_batch, n_features)`.
> 
> My understanding is that the goal of this PR is to implement the above functionality. I am wondering whether your implementation @jackaraz is `TensorFlow` specific. Have you tested whether it can be used without an `interface`, i.e. vanilla `pennylane`, or with `PyTorch` (`qnn.TorchLayer`)?

Hi @vbelis, exactly the goal is to batch like in ML context. The function that I implemented should be generic, I tested with pennylane NumPy interface and TF but not PyTorch since I don't use PyTorch. However, since it uses the penny lane backend the separation between trainable and non-trainable parameters should be already available to the `batch_transform` class hence the function should be able to work with PyTorch as well. 
vbelis: Hi @jackaraz, thanks for the fast reply. OK I will pull and do some tests locally to check the behavior, since I am using `PyTorch`. The PR has not been concluded due to the design mismatch that the CI flagged or are you also trying to enhance it further before pushing?
jackaraz: Hi @vbelis I wasn't planning to add anything else, not really sure why the tests are failing, if you have any feedback would be much appreciated to finalize this PR. 
vbelis: I will try to take a look soon. However, it would be much more efficient if some of the main developers, e.g. @josh146, could provide some tips regarding the failed tests: (`codecov` and `sphinx`). In the `CodeFactor` failed check, they seem to have an upper limit of 5 arguments for the constructor. Do you believe there is a good way to remove that argument or include it in the keyword arguments, to stay consistent with their code design? Might be good to also merge the master branch to your branch to get the latest commits (there is some warning about this as well).

Regarding the failed checks by `codecov/patch` of the form: `Added line #L... was not covered by tests`, I don't have any suggestion out of the top of my head, sorry... 
josh146: > Regarding the failed checks by codecov/patch of the form: Added line #L... was not covered by tests, I don't have any suggestion out of the top of my head, sorry...

This is a CI check that verifies that all lines of code added in this PR are tested :) In this case, it appears that this particular line is not being called by any of the unit tests.

>  OK I will pull and do some tests locally to check the behavior, since I am using PyTorch. 

This is much appreciated @vbelis! Let me know if you have any questions :)

jackaraz: @vbelis just updated with the batch input test and added the line @josh146 suggested. Would you like to add batching in PyTorch as well? 
bb511: Hey @jackaraz, I am a collaborator of @vbelis. We are currently testing the pytroch functionality to see if we need to add additional lines to the code. We'll let you know if no further changes to the code are required. Otherwise, we'll push these additional changes to this PR since we think that creating another one would be superfluous.
jackaraz: > Hey @jackaraz, I am a collaborator of @vbelis. We are currently testing the pytroch functionality to see if we need to add additional lines to the code. We'll let you know if no further changes to the code are required. Otherwise, we'll push these additional changes to this PR since we think that creating another one would be superfluous.

Hi @bb511 I totally agree with you, please do so. Note that I haven't added anything to pytorch module but in principle, it should be able to use this batching function as TensorFlow does since it's using purely pennylane backend for batching.
jackaraz: @bb511 @vbelis PS you can find more comprehensive and independent tests/executions in [here](https://gitlab.com/jackaraz/tnqcircuits) as well.
josh146: @jackaraz let me have a look 🙂 

It looks like only the documentation check is failing, indicating an issue with the docstrings.

If you click on the 'details' next to the failing check, and search 'warnings', you will see this in the log:

```
/github/workspace/pennylane/qnn/keras.py:docstring of pennylane.qnn.keras.KerasLayer:25: WARNING: Unexpected indentation.
```

I'll have a look at the docstring and see if I can spot the incorrect indentation :)
jackaraz: Thanks @josh146; all suggestions have been implemented.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index e471892bb..16f88892c 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -296,6 +296,9 @@
     device.
     [(#2257)](https://github.com/PennyLaneAI/pennylane/pull/2257)
 
+* 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>
 
 * Some error messages have been improved and operation properties made
@@ -547,6 +550,6 @@ The Operator class has undergone a major refactor with the following changes:
 
 This release contains contributions from (in alphabetical order):
 
-Sam Banning, Thomas Bromley, Anthony Hayes, David Ittah, Josh Izaac, Christina Lee, Angus Lowe,
-Maria Fernanda Morris, Romain Moyard, Zeyue Niu, Maria Schuld, Jay Soni,
+Jack Y. Araz, Sam Banning, Thomas Bromley, Olivia Di Matteo, Anthony Hayes, David Ittah, Josh Izaac,
+Christina Lee, Angus Lowe, Maria Fernanda Morris, Romain Moyard, Zeyue Niu, Maria Schuld, Jay Soni,
 Antal Száva, David Wierichs
diff --git a/pennylane/__init__.py b/pennylane/__init__.py
index 2df4026b4..363d90484 100644
--- a/pennylane/__init__.py
+++ b/pennylane/__init__.py
@@ -56,6 +56,7 @@ from pennylane.transforms import (
     adjoint,
     adjoint_metric_tensor,
     batch_params,
+    batch_input,
     batch_transform,
     draw,
     draw_mpl,
diff --git a/pennylane/measurements.py b/pennylane/measurements.py
index fa3007505..c4358b0a0 100644
--- a/pennylane/measurements.py
+++ b/pennylane/measurements.py
@@ -17,6 +17,7 @@ This module contains the functions for computing different types of measurement
 outcomes from quantum observables - expectation values, variances of expectations,
 and measurement samples using AnnotatedQueues.
 """
+# pylint: disable=too-many-instance-attributes
 import copy
 import uuid
 import warnings
diff --git a/pennylane/qnn/keras.py b/pennylane/qnn/keras.py
index 582813127..26e5565cf 100644
--- a/pennylane/qnn/keras.py
+++ b/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 :func:`~.pennylane.transforms.batch_input` for more details.
         **kwargs: additional keyword arguments passed to the Layer_ base class
 
     **Example**
@@ -196,8 +202,15 @@ class KerasLayer(Layer):
     """
 
     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 @@ class KerasLayer(Layer):
         }
 
         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 @@ class KerasLayer(Layer):
         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 @@ class KerasLayer(Layer):
         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 @@ class KerasLayer(Layer):
         `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
diff --git a/pennylane/transforms/__init__.py b/pennylane/transforms/__init__.py
index d4775b22c..64583d258 100644
--- a/pennylane/transforms/__init__.py
+++ b/pennylane/transforms/__init__.py
@@ -31,6 +31,7 @@ that compute the desired quantity.
 
     ~transforms.classical_jacobian
     ~batch_params
+    ~batch_input
     ~draw
     ~draw_mpl
     ~transforms.get_unitary_matrix
@@ -161,6 +162,7 @@ from .qfunc_transforms import make_tape, single_tape_transform, qfunc_transform
 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 .condition import cond, Conditional
 from .compile import compile
diff --git a/pennylane/transforms/batch_input.py b/pennylane/transforms/batch_input.py
new file mode 100644
index 000000000..da48b8824
--- /dev/null
+++ b/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.QuantumTape 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.QuantumTape], Callable]: list of tapes arranged
+        according to unbatched inputs and a callable function to batch the results.
+
+    .. seealso:: :func:`~.batch_params`
+
+    **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])>
+    """
+    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))


Test patch:
diff --git a/tests/qnn/test_keras.py b/tests/qnn/test_keras.py
index f26895ce1..4c123be90 100644
--- a/tests/qnn/test_keras.py
+++ b/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)
diff --git a/tests/transforms/test_batch_inputs.py b/tests/transforms/test_batch_inputs.py
new file mode 100644
index 000000000..c03fd3198
--- /dev/null
+++ b/tests/transforms/test_batch_inputs.py
@@ -0,0 +1,128 @@
+# 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 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, 1])
+    @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,))
+
+    with pytest.raises(ValueError, match="Batch dimension for all gate arguments"):
+        res = circuit(input1, input2, weights)
+
+
+def test_mottonenstate_preparation(mocker):
+    """Test that batching works for MottonenStatePreparation"""
+    dev = qml.device("default.qubit", wires=3)
+
+    @qml.batch_input(argnum=0)
+    @qml.qnode(dev)
+    def circuit(data, weights):
+        qml.templates.MottonenStatePreparation(data, wires=[0, 1, 2])
+        qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1, 2])
+        return qml.probs(wires=[0, 1, 2])
+
+    batch_size = 3
+
+    # create a batched input statevector
+    data = np.random.random((batch_size, 2**3), requires_grad=False)
+    data /= np.linalg.norm(data, axis=1).reshape(-1, 1)  # normalize
+
+    # weights is not batched
+    weights = np.random.random((10, 3, 3), requires_grad=True)
+
+    spy = mocker.spy(circuit.device, "batch_execute")
+    res = circuit(data, weights)
+    assert res.shape == (batch_size, 2**3)
+    assert len(spy.call_args[0][0]) == batch_size
+
+    # check the results against individually executed circuits (no batching)
+    @qml.qnode(dev)
+    def circuit2(data, weights):
+        qml.templates.MottonenStatePreparation(data, wires=[0, 1, 2])
+        qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1, 2])
+        return qml.probs(wires=[0, 1, 2])
+
+    indiv_res = []
+    for state in data:
+        indiv_res.append(circuit2(state, weights))
+    assert np.allclose(res, indiv_res)
+
+
+@pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift"])
+def test_autograd(diff_method, tol):
+    """Test derivatives when using autograd"""
+    dev = qml.device("default.qubit", wires=2)
+
+    @qml.batch_input
+    @qml.qnode(dev, diff_method=diff_method)
+    def circuit(x):
+        qml.RX(x, wires=0)
+        qml.RY(0.1, wires=1)
+        qml.CNOT(wires=[0, 1])
+        return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+
+    def cost(x):
+        return np.sum(circuit(x))
+
+    batch_size = 3
+    x = np.linspace(0.1, 0.5, batch_size, requires_grad=True)
+
+    res = qml.grad(cost)(x)
+    expected = -np.sin(0.1) * np.sin(x)
+    assert np.allclose(res, expected, atol=tol, rtol=0)