Repository: PennyLaneAI/pennylane

Issue:
[BUG] QAOA layer template does not work with TensorFlow
### Expected behavior

PennyLane's QAOA layer templates should work with the TensorFlow interface.

### Actual behavior

A `TypeError: object of type 'ResourceVariable' has no len()` is thrown.

### Additional information

_No response_

### Source code

```shell
import networkx as nx
import pennylane as qml
import numpy as np
import tensorflow as tf

num_nodes = 4
num_edges = 4
p = 2
seed = 1967

def qaoa_layer(gamma, alpha):
    qml.qaoa.cost_layer(gamma, cost_h)
    qml.qaoa.mixer_layer(alpha, mixer_h)

def circuit(gamma, alpha, **kwargs):
    for i in range(num_nodes):
        qml.Hadamard(wires=i)
        
    qml.layer(qaoa_layer, p, gamma, alpha)

g = nx.gnm_random_graph(num_nodes, num_edges, seed=seed)
cost_h, mixer_h = qml.qaoa.maxcut(g)

dev = qml.device('default.qubit', wires=4)

alpha = tf.Variable(np.random.rand(p), dtype=tf.float64)
gamma = tf.Variable(np.random.rand(p), dtype=tf.float64)
cost_function = qml.ExpvalCost(circuit, cost_h, dev, interface='tf', optimize=True)
cost_function(gamma, alpha)
```


### Tracebacks

```shell
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
/var/folders/v1/y849xjw53gl8_4b7cf9qr4pxc7s0j7/T/ipykernel_1781/1824392795.py in <module>
----> 1 cost_function(gamma, alpha)

~/.pyenv/versions/3.9.6/envs/braket-dev-3.9.6/lib/python3.9/site-packages/pennylane/vqe/vqe.py in __call__(self, *args, **kwargs)
    585 
    586     def __call__(self, *args, **kwargs):
--> 587         return self.cost_fn(*args, **kwargs)
    588 
    589 

~/.pyenv/versions/3.9.6/envs/braket-dev-3.9.6/lib/python3.9/site-packages/pennylane/vqe/vqe.py in cost_fn(*qnode_args, **qnode_kwargs)
    575                 total = 0
    576                 for o, c in zip(obs_groupings, coeffs_groupings):
--> 577                     res = circuit(*qnode_args, obs=o, **qnode_kwargs)
    578                     total += sum([r * c_ for r, c_ in zip(res, c)])
    579                 return total

~/.pyenv/versions/3.9.6/envs/braket-dev-3.9.6/lib/python3.9/site-packages/pennylane/qnode.py in __call__(self, *args, **kwargs)
    593         if self.mutable or self.qtape is None:
    594             # construct the tape
--> 595             self.construct(args, kwargs)
    596 
    597         # execute the tape

~/.pyenv/versions/3.9.6/envs/braket-dev-3.9.6/lib/python3.9/site-packages/pennylane/qnode.py in construct(self, args, kwargs)
    522 
    523         with self.qtape:
--> 524             self.qfunc_output = self.func(*args, **kwargs)
    525 
    526         if not isinstance(self.qfunc_output, Sequence):

~/.pyenv/versions/3.9.6/envs/braket-dev-3.9.6/lib/python3.9/site-packages/pennylane/vqe/vqe.py in circuit(obs, *qnode_args, **qnode_kwargs)
    568             def circuit(*qnode_args, obs, **qnode_kwargs):
    569                 """Converting ansatz into a full circuit including measurements"""
--> 570                 ansatz(*qnode_args, wires=w, **qnode_kwargs)
    571                 return [qml.expval(o) for o in obs]
    572 

/var/folders/v1/y849xjw53gl8_4b7cf9qr4pxc7s0j7/T/ipykernel_1781/926742614.py in circuit(gamma, alpha, **kwargs)
      3         qml.Hadamard(wires=i)
      4 
----> 5     qml.layer(qaoa_layer, p, gamma, alpha)

~/.pyenv/versions/3.9.6/envs/braket-dev-3.9.6/lib/python3.9/site-packages/pennylane/templates/decorator.py in wrapper(*args, **kwargs)
     59 
     60         with qml.tape.OperationRecorder() as rec:
---> 61             func(*args, **kwargs)
     62 
     63         return rec.queue

~/.pyenv/versions/3.9.6/envs/braket-dev-3.9.6/lib/python3.9/site-packages/pennylane/templates/layer.py in layer(template, depth, *args, **kwargs)
    211     """
    212 
--> 213     _preprocess(args, depth)
    214 
    215     for i in range(0, int(depth)):

~/.pyenv/versions/3.9.6/envs/braket-dev-3.9.6/lib/python3.9/site-packages/pennylane/templates/layer.py in _preprocess(args, depth)
     31     for arg in args:
     32 
---> 33         if len(arg) != depth:
     34             raise ValueError(
     35                 "Each positional argument must have length matching 'depth'; expected {} got {}".format(

TypeError: object of type 'ResourceVariable' has no len()
```


### System information

```shell
Name: PennyLane
Version: 0.16.0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/XanaduAI/pennylane
Author: None
Author-email: None
License: Apache License 2.0
Location: /Users/licedric/.pyenv/versions/3.9.6/envs/braket-dev-3.9.6/lib/python3.9/site-packages
Requires: scipy, autoray, autograd, appdirs, networkx, toml, semantic-version, numpy
Required-by: PennyLane-Qchem, amazon-braket-pennylane-plugin
Platform info:           macOS-10.15.7-x86_64-i386-64bit
Python version:          3.9.6
Numpy version:           1.19.5
Scipy version:           1.7.0
Installed devices:
- default.gaussian (PennyLane-0.16.0)
- default.mixed (PennyLane-0.16.0)
- default.qubit (PennyLane-0.16.0)
- default.qubit.autograd (PennyLane-0.16.0)
- default.qubit.jax (PennyLane-0.16.0)
- default.qubit.tf (PennyLane-0.16.0)
- default.tensor (PennyLane-0.16.0)
- default.tensor.tf (PennyLane-0.16.0)
- braket.aws.qubit (amazon-braket-pennylane-plugin-1.4.0)
- braket.local.qubit (amazon-braket-pennylane-plugin-1.4.0)
```


### 

- [X] I have searched exisisting GitHub issues to make sure the issue does not already exist.


Discussion:
Thanks @licedric! From looking into this further, it appears that the issue is that `qml.layer` uses `len(args)`, which will not work if any of the input arguments to be layered are `tf.Variable` objects which are not iterable.

A solution would be to either:

- Avoid using `len()`, by replacing it with `qml.math.shape(args)[0]`
- Convert `args` to a `tf.Tensor` (which _is_ iterable) if it happens to be a `tf,Variable`.

The first option is preferable, since it is framework agnostic.

We should also add TF and Torch tests for `qml.layer`.

Resolving patch:
diff --git a/pennylane/templates/layer.py b/pennylane/templates/layer.py
--- a/pennylane/templates/layer.py
+++ b/pennylane/templates/layer.py
@@ -16,6 +16,7 @@
 """
 # pylint: disable-msg=too-many-branches,too-many-arguments,protected-access
 from pennylane.templates.decorator import template as temp
+from pennylane.math import shape
 
 
 def _preprocess(args, depth):
@@ -30,7 +31,7 @@ def _preprocess(args, depth):
 
     for arg in args:
 
-        if len(arg) != depth:
+        if shape(arg)[0] != depth:
             raise ValueError(
                 "Each positional argument must have length matching 'depth'; expected {} got {}".format(
                     depth, len(arg)


Test patch:
diff --git a/tests/templates/test_layer.py b/tests/templates/test_layer.py
--- a/tests/templates/test_layer.py
+++ b/tests/templates/test_layer.py
@@ -155,3 +155,22 @@ def test_layer(self, unitary, depth, arguments, keywords, gates):
             target = [gates[i].name, gates[i].parameters, gates[i].wires]
 
         assert prep == target
+
+    def test_layer_tf(self):
+        """Tests that the layering function accepts Tensorflow variables."""
+
+        tf = pytest.importorskip("tensorflow")
+
+        def unitary(param):
+            qml.RX(param, wires=0)
+
+        x = tf.Variable([0.1, 0.2, 0.3])
+
+        with qml.tape.OperationRecorder() as rec:
+            layer(unitary, 3, x)
+
+        assert len(rec.operations) == 3
+
+        for ii, op in enumerate(rec.operations):
+            assert qml.math.allclose(op.parameters[0], x[ii])
+            assert isinstance(op, qml.RX)