Repository: PennyLaneAI/pennylane

Issue:
Allow parameter broadcasting in `AmplitudeEmbedding` with TensorFlow
### Feature details

As of #2810, embedding templates support parameter broadcasting, including `AmplitudeEmbedding` except for usage with Tensorflow, which is raising a dedicated error currently.
To reach feature parity, one would have to figure out how to pad the parameter-broadcasted amplitudes and normalize them correctly across the non-broadcasting dimension.

### Implementation

_No response_

### How important would you say this feature is?

1: Not important. Would be nice to have.

### Additional information

_No response_


Discussion:
We received another request for this feature here: https://discuss.pennylane.ai/t/importerror-cannot-import-name-shape/1383/17?u=isaacdevlugt

Let's make sure to post any updates on there as well 😄 

Resolving patch:
diff --git a/pennylane/templates/embeddings/amplitude.py b/pennylane/templates/embeddings/amplitude.py
--- a/pennylane/templates/embeddings/amplitude.py
+++ b/pennylane/templates/embeddings/amplitude.py
@@ -15,8 +15,6 @@
 Contains the AmplitudeEmbedding template.
 """
 # pylint: disable-msg=too-many-branches,too-many-arguments,protected-access
-import numpy as np
-
 import pennylane as qml
 from pennylane.ops import StatePrep
 from pennylane.wires import Wires
@@ -167,63 +165,51 @@ def _preprocess(features, wires, pad_with, normalize):
         * If normalize is false, check that last dimension of features is normalised to one. Else, normalise the
           features tensor.
         """
+        shape = qml.math.shape(features)
+
+        # check shape
+        if len(shape) not in (1, 2):
+            raise ValueError(
+                f"Features must be a one-dimensional tensor, or two-dimensional with batching; got shape {shape}."
+            )
+
+        n_features = shape[-1]
+        dim = 2 ** len(wires)
+        if pad_with is None and n_features != dim:
+            raise ValueError(
+                f"Features must be of length {dim}; got length {n_features}. "
+                f"Use the 'pad_with' argument for automated padding."
+            )
+
+        if pad_with is not None:
+            if n_features > dim:
+                raise ValueError(
+                    f"Features must be of length {dim} or "
+                    f"smaller to be padded; got length {n_features}."
+                )
 
-        # check if features is batched
-        batched = qml.math.ndim(features) > 1
-
-        if batched and qml.math.get_interface(features) == "tensorflow":
-            raise ValueError("AmplitudeEmbedding does not support batched Tensorflow features.")
-
-        features_batch = features if batched else [features]
-
-        new_features_batch = []
-        # apply pre-processing to each features tensor in the batch
-        for feature_set in features_batch:
-            shape = qml.math.shape(feature_set)
-
-            # check shape
-            if len(shape) != 1:
-                raise ValueError(f"Features must be a one-dimensional tensor; got shape {shape}.")
-
-            n_features = shape[0]
-            dim = 2 ** len(wires)
-            if pad_with is None and n_features != dim:
+            # pad
+            if n_features < dim:
+                padding = [pad_with] * (dim - n_features)
+                if len(shape) > 1:
+                    padding = [padding] * shape[0]
+                padding = qml.math.convert_like(padding, features)
+                features = qml.math.hstack([features, padding])
+
+        # normalize
+        norm = qml.math.sum(qml.math.abs(features) ** 2, axis=-1)
+
+        if qml.math.is_abstract(norm):
+            if normalize or pad_with:
+                features = features / qml.math.reshape(qml.math.sqrt(norm), (*shape[:-1], 1))
+
+        elif not qml.math.allclose(norm, 1.0, atol=TOLERANCE):
+            if normalize or pad_with:
+                features = features / qml.math.reshape(qml.math.sqrt(norm), (*shape[:-1], 1))
+            else:
                 raise ValueError(
-                    f"Features must be of length {dim}; got length {n_features}. "
-                    f"Use the 'pad_with' argument for automated padding."
+                    f"Features must be a vector of norm 1.0; got norm {norm}. "
+                    "Use 'normalize=True' to automatically normalize."
                 )
 
-            if pad_with is not None:
-                if n_features > dim:
-                    raise ValueError(
-                        f"Features must be of length {dim} or "
-                        f"smaller to be padded; got length {n_features}."
-                    )
-
-                # pad
-                if n_features < dim:
-                    padding = [pad_with] * (dim - n_features)
-                    padding = qml.math.convert_like(padding, feature_set)
-                    feature_set = qml.math.hstack([feature_set, padding])
-
-            # normalize
-            norm = qml.math.sum(qml.math.abs(feature_set) ** 2)
-
-            if qml.math.is_abstract(norm):
-                if normalize or pad_with:
-                    feature_set = feature_set / qml.math.sqrt(norm)
-
-            elif not qml.math.allclose(norm, 1.0, atol=TOLERANCE):
-                if normalize or pad_with:
-                    feature_set = feature_set / qml.math.sqrt(norm)
-                else:
-                    raise ValueError(
-                        f"Features must be a vector of norm 1.0; got norm {norm}. "
-                        "Use 'normalize=True' to automatically normalize."
-                    )
-
-            new_features_batch.append(feature_set)
-
-        return qml.math.cast(
-            qml.math.stack(new_features_batch) if batched else new_features_batch[0], np.complex128
-        )
+        return features


Test patch:
diff --git a/tests/templates/test_embeddings/test_amplitude.py b/tests/templates/test_embeddings/test_amplitude.py
--- a/tests/templates/test_embeddings/test_amplitude.py
+++ b/tests/templates/test_embeddings/test_amplitude.py
@@ -286,14 +286,27 @@ def test_id(self):
         assert template.id == "a"
 
 
-def circuit_template(features):
-    qml.AmplitudeEmbedding(features, wires=range(3))
+def circuit_template(features, pad_with=None, normalize=False):
+    """AmplitudeEmbedding circuit. For three wires, all test features match
+    the expected length. For four wires, we get a test case for `pad_with`."""
+    num_wires = 3 if pad_with is None else 4
+    qml.AmplitudeEmbedding(features, wires=range(num_wires), pad_with=pad_with, normalize=normalize)
     return qml.state()
 
 
-def circuit_decomposed(features):
+def circuit_decomposed(features, pad_with=None):
+    """AmplitudeEmbedding circuit reexpressed as manual state preparation.
+    This function expects the length of the features to match."""
+    num_wires = 3 if pad_with is None else 4
+    if pad_with is not None:
+        # and need to pad manually in order to double the size of the vector
+        # from 8 (3 qubits) to 16 (4 qubits). Also, normalize
+        features = qml.math.hstack([features, qml.math.ones_like(features) * pad_with])
+        shape = qml.math.shape(features)
+        norm = qml.math.reshape(qml.math.linalg.norm(features, axis=-1), (*shape[:-1], 1))
+        features = features / norm
     # need to cast to complex tensor, which is implicitly done in the template
-    qml.StatePrep(qml.math.cast(features, np.complex128), wires=range(3))
+    qml.StatePrep(qml.math.cast(features, np.complex128), wires=range(num_wires))
     return qml.state()
 
 
@@ -307,267 +320,128 @@ def circuit_decomposed(features):
 ]
 
 
+@pytest.mark.parametrize("features", all_features)
+@pytest.mark.parametrize("pad_with", [None, 0.0, 0.1])
+@pytest.mark.parametrize("normalize", [False, True])
 class TestInterfaces:
     """Tests that the template is compatible with all interfaces."""
 
-    @pytest.mark.parametrize("features", all_features)
-    def test_list_and_tuples(self, tol, features):
+    def test_list_and_tuples(self, tol, features, pad_with, normalize):
         """Tests common iterables as inputs."""
 
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit")
 
         circuit = qml.QNode(circuit_template, dev)
         circuit2 = qml.QNode(circuit_decomposed, dev)
 
-        res = circuit(features)
-        res2 = circuit2(features)
+        res = circuit(features, pad_with, normalize)
+        res2 = circuit2(features, pad_with)
         assert qml.math.allclose(res, res2, atol=tol, rtol=0)
 
-        res = circuit(tuple(features))
-        res2 = circuit2(tuple(features))
+        res = circuit(tuple(features), pad_with, normalize)
+        res2 = circuit2(tuple(features), pad_with)
         assert qml.math.allclose(res, res2, atol=tol, rtol=0)
 
     @pytest.mark.autograd
-    @pytest.mark.parametrize("features", all_features)
-    def test_autograd(self, tol, features):
+    def test_autograd(self, tol, features, pad_with, normalize):
         """Tests autograd tensors."""
 
         features = pnp.array(features, requires_grad=True)
 
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit")
 
-        circuit = qml.QNode(circuit_template, dev)
+        circuit = qml.QNode(circuit_template, dev, interface="autograd")
         circuit2 = qml.QNode(circuit_decomposed, dev)
 
-        res = circuit(features)
-        res2 = circuit2(features)
-
-        assert qml.math.allclose(res, res2, atol=tol, rtol=0)
-
-    @pytest.mark.autograd
-    @pytest.mark.parametrize("features", all_features)
-    def test_autograd_pad_with(self, tol, features):
-        """Tests autograd tensors and pad_with."""
-
-        features = pnp.array(features)
-
-        dev = qml.device("default.qubit", wires=4)
-
-        @qml.qnode(dev)
-        def node_decomposed(features):
-            # need to cast to complex tensor, which is implicitly done in the template
-            state = qml.math.cast(
-                qml.math.hstack([features, qml.math.zeros_like(features)]), np.complex128
-            )
-            qml.StatePrep(state, wires=range(4))
-            return qml.state()
-
-        @qml.qnode(dev)
-        def circuit(x=None):
-            qml.AmplitudeEmbedding(x, list(range(4)), pad_with=0.0)
-            return qml.state()
-
-        res = circuit(features)
-        res2 = node_decomposed(features)
+        res = circuit(features, pad_with, normalize)
+        res2 = circuit2(features, pad_with)
 
         assert qml.math.allclose(res, res2, atol=tol, rtol=0)
 
     @pytest.mark.jax
-    @pytest.mark.parametrize("features", all_features)
-    def test_jax(self, tol, features):
+    def test_jax(self, tol, features, pad_with, normalize):
         """Tests jax tensors."""
-
         import jax.numpy as jnp
 
         features = jnp.array(features)
 
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit")
 
-        circuit = qml.QNode(circuit_template, dev)
+        circuit = qml.QNode(circuit_template, dev, interface="jax")
         circuit2 = qml.QNode(circuit_decomposed, dev)
 
-        res = circuit(features)
-        res2 = circuit2(features)
+        res = circuit(features, pad_with, normalize)
+        res2 = circuit2(features, pad_with)
 
         assert qml.math.allclose(res, res2, atol=tol, rtol=0)
 
     @pytest.mark.jax
-    @pytest.mark.parametrize("features", all_features)
-    def test_jax_jit(self, tol, features):
-        """Tests jax tensors when using JIT."""
-
+    def test_jax_jit(self, tol, features, pad_with, normalize):
+        """Tests jax tensors with JIT compilation."""
         import jax
         import jax.numpy as jnp
 
         features = jnp.array(features)
 
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit")
 
-        circuit = jax.jit(qml.QNode(circuit_template, dev))
-        circuit2 = jax.jit(qml.QNode(circuit_decomposed, dev))
+        circuit = jax.jit(qml.QNode(circuit_template, dev, interface="jax"), static_argnums=[1, 2])
+        circuit2 = jax.jit(qml.QNode(circuit_decomposed, dev), static_argnums=[1])
 
-        res = circuit(features)
-        res2 = circuit2(features)
-
-        assert qml.math.allclose(res, res2, atol=tol, rtol=0)
-
-    @pytest.mark.jax
-    @pytest.mark.parametrize("features", all_features)
-    def test_jax_jit_pad_with(self, tol, features):
-        """Tests jax tensors when using JIT and pad_with."""
-
-        import jax
-        import jax.numpy as jnp
-
-        features = jnp.array(features)
-
-        dev = qml.device("default.qubit", wires=4)
-
-        @jax.jit
-        @qml.qnode(dev)
-        def node_decomposed(features):
-            # need to cast to complex tensor, which is implicitly done in the template
-            state = qml.math.cast(
-                qml.math.hstack([features, qml.math.zeros_like(features)]), np.complex128
-            )
-            qml.StatePrep(state, wires=range(4))
-            return qml.state()
-
-        @jax.jit
-        @qml.qnode(dev)
-        def circuit(x=None):
-            qml.AmplitudeEmbedding(x, list(range(4)), pad_with=0.0)
-            return qml.state()
-
-        res = circuit(features)
-        res2 = node_decomposed(features)
+        res = circuit(features, pad_with, normalize)
+        res2 = circuit2(features, pad_with)
 
         assert qml.math.allclose(res, res2, atol=tol, rtol=0)
 
     @pytest.mark.tf
-    def test_tf(self, tol):
-        """Tests tf tensors."""
-
+    def test_tf(self, tol, features, pad_with, normalize):
+        """Tests tensorflow tensors."""
         import tensorflow as tf
 
-        features = tf.Variable(all_features[0])
+        features = tf.Variable(features)
 
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit")
 
-        circuit = qml.QNode(circuit_template, dev)
+        circuit = qml.QNode(circuit_template, dev, interface="tensorflow")
         circuit2 = qml.QNode(circuit_decomposed, dev)
 
-        res = circuit(features)
-        res2 = circuit2(features)
-
-        assert qml.math.allclose(res, res2, atol=tol, rtol=0)
-
-    @pytest.mark.tf
-    @pytest.mark.parametrize("features", all_features)
-    def test_tf_pad_with(self, tol, features):
-        """Tests tf tensors and pad_with."""
-        import tensorflow as tf
-
-        features = tf.Variable(all_features[0])
-
-        dev = qml.device("default.qubit", wires=4)
-
-        @qml.qnode(dev)
-        def node_decomposed(features):
-            # need to cast to complex tensor, which is implicitly done in the template
-            state = qml.math.cast(
-                qml.math.hstack([features, qml.math.zeros_like(features)]), tf.complex128
-            )
-            qml.StatePrep(state, wires=range(4))
-            return qml.state()
-
-        @qml.qnode(dev)
-        def circuit(x=None):
-            qml.AmplitudeEmbedding(x, list(range(4)), pad_with=0.0)
-            return qml.state()
-
-        res = circuit(features)
-        res2 = node_decomposed(features)
+        res = circuit(features, pad_with, normalize)
+        res2 = circuit2(features, pad_with)
 
         assert qml.math.allclose(res, res2, atol=tol, rtol=0)
 
     @pytest.mark.tf
-    def test_tf_error_when_batching(self):
-        """Tests batched tf tensors raising an error."""
-
+    def test_tf_jit(self, tol, features, pad_with, normalize):
+        """Tests tensorflow tensors with JIT compilation."""
         import tensorflow as tf
 
-        features = tf.Variable(all_features[1])
-
-        dev = qml.device("default.qubit", wires=3)
-
-        circuit = qml.QNode(circuit_template, dev)
-
-        with pytest.raises(ValueError, match="does not support batched Tensorflow features"):
-            circuit(features)
-
-    @pytest.mark.tf
-    def test_tf_jit(self, tol):
-        """Tests tf tensors when using JIT."""
-
-        import tensorflow as tf
+        features = tf.Variable(features)
 
-        features = tf.Variable(all_features[0])
+        dev = qml.device("default.qubit")
 
-        dev = qml.device("default.qubit", wires=3)
-
-        circuit = tf.function(jit_compile=True)(qml.QNode(circuit_template, dev))
+        circuit = tf.function(jit_compile=True)(
+            qml.QNode(circuit_template, dev, interface="tensorflow")
+        )
         circuit2 = tf.function(jit_compile=True)(qml.QNode(circuit_decomposed, dev))
 
-        res = circuit(features)
-        res2 = circuit2(features)
+        res = circuit(features, pad_with, normalize)
+        res2 = circuit2(features, pad_with)
 
         assert qml.math.allclose(res, res2, atol=tol, rtol=0)
 
     @pytest.mark.torch
-    @pytest.mark.parametrize("features", all_features)
-    def test_torch(self, tol, features):
-        """Tests torch tensors."""
-
+    def test_torch(self, tol, features, pad_with, normalize):
+        """Tests Torch tensors."""
         import torch
 
         features = torch.tensor(features, requires_grad=True)
 
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit")
 
-        circuit = qml.QNode(circuit_template, dev)
+        circuit = qml.QNode(circuit_template, dev, interface="torch")
         circuit2 = qml.QNode(circuit_decomposed, dev)
 
-        res = circuit(features)
-        res2 = circuit2(features)
-
-        assert qml.math.allclose(res, res2, atol=tol, rtol=0)
-
-    @pytest.mark.torch
-    @pytest.mark.parametrize("features", all_features)
-    def test_torch_pad_with(self, tol, features):
-        """Tests torch tensors and pad_with."""
-        import torch
-
-        features = torch.tensor(features, requires_grad=True)
-
-        dev = qml.device("default.qubit", wires=4)
-
-        @qml.qnode(dev)
-        def node_decomposed(features):
-            # need to cast to complex tensor, which is implicitly done in the template
-            state = qml.math.cast(
-                qml.math.hstack([features, qml.math.zeros_like(features)]), torch.complex128
-            )
-            qml.StatePrep(state, wires=range(4))
-            return qml.state()
-
-        @qml.qnode(dev)
-        def circuit(x=None):
-            qml.AmplitudeEmbedding(x, list(range(4)), pad_with=0.0)
-            return qml.state()
-
-        res = circuit(features)
-        res2 = node_decomposed(features)
+        res = circuit(features, pad_with, normalize)
+        res2 = circuit2(features, pad_with)
 
         assert qml.math.allclose(res, res2, atol=tol, rtol=0)