Repository: PennyLaneAI/pennylane
Pull request #5969 (2024-07-08 21:21:41)

Fix `AmplitudeEmbedding` to handle integer dtypes
**Context:**
Fixes bug #5940 .

**Description of the Change:**
* Use `qml.math.linalg.norm` instead of manually taking the sum of squared features for normalizing. To get this to work, I also  cast to floats whenever `features` is an integer dtype because tf and torch don't allow using `norm` on integer arrays.

**Benefits:**
* `qml.AmplitudeEmbedding` works with all dtypes.

**Possible Drawbacks:**

**Related GitHub Issues:**

[sc-67672]

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index f331c2273..ff8f94d2b 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -34,6 +34,9 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* `qml.AmplitudeEmbedding` has better support for features using low precision integer data types.
+[(#5969)](https://github.com/PennyLaneAI/pennylane/pull/5969)
+
 <h3>Contributors ✍️</h3>
 
 This release contains contributions from (in alphabetical order):
@@ -41,4 +44,5 @@ Yushao Chen,
 Christina Lee,
 William Maxwell,
 Vincent Michaud-Rioux,
+Mudit Pandey,
 Erik Schultheis.
diff --git a/pennylane/templates/embeddings/amplitude.py b/pennylane/templates/embeddings/amplitude.py
index f87efbbec..7fb139dc1 100644
--- a/pennylane/templates/embeddings/amplitude.py
+++ b/pennylane/templates/embeddings/amplitude.py
@@ -165,6 +165,8 @@ class AmplitudeEmbedding(StatePrep):
         * If normalize is false, check that last dimension of features is normalised to one. Else, normalise the
           features tensor.
         """
+        if isinstance(features, (list, tuple)):
+            features = qml.math.array(features)
         shape = qml.math.shape(features)
 
         # check shape
@@ -197,15 +199,17 @@ class AmplitudeEmbedding(StatePrep):
                 features = qml.math.hstack([features, padding])
 
         # normalize
-        norm = qml.math.sum(qml.math.abs(features) ** 2, axis=-1)
+        if "int" in str(features.dtype):
+            features = qml.math.cast_like(features, 0.0)
+        norm = qml.math.linalg.norm(features, 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))
+                features = features / qml.math.reshape(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))
+                features = features / qml.math.reshape(norm, (*shape[:-1], 1))
             else:
                 raise ValueError(
                     f"Features must be a vector of norm 1.0; got norm {norm}. "


Test patch:
diff --git a/tests/templates/test_embeddings/test_amplitude.py b/tests/templates/test_embeddings/test_amplitude.py
index bab564a6d..5cf0f47c8 100644
--- a/tests/templates/test_embeddings/test_amplitude.py
+++ b/tests/templates/test_embeddings/test_amplitude.py
@@ -302,14 +302,18 @@ def circuit_template(features, pad_with=None, normalize=False):
     return qml.state()
 
 
-def circuit_decomposed(features, pad_with=None):
+def circuit_decomposed(features, pad_with=None, normalize=False):
     """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 not isinstance(features, (list, tuple)):
+        features = qml.math.cast_like(features, 1.0) if "int" in str(features.dtype) else features
     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
+        pad_with = qml.math.cast_like(pad_with, features)
         features = qml.math.hstack([features, qml.math.ones_like(features) * pad_with])
+    if pad_with is not None or normalize:
         shape = qml.math.shape(features)
         norm = qml.math.reshape(qml.math.linalg.norm(features, axis=-1), (*shape[:-1], 1))
         features = features / norm
@@ -394,7 +398,7 @@ class TestInterfaces:
         dev = qml.device("default.qubit")
 
         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])
+        circuit2 = jax.jit(qml.QNode(circuit_decomposed, dev), static_argnums=[1, 2])
 
         res = circuit(features, pad_with, normalize)
         res2 = circuit2(features, pad_with)
@@ -453,3 +457,130 @@ class TestInterfaces:
         res2 = circuit2(features, pad_with)
 
         assert qml.math.allclose(res, res2, atol=tol, rtol=0)
+
+
+int_features = [
+    [10, 0, 2, 35, 0, 41, 56, 0],
+    [
+        [10, 0, 2, 35, 0, 41, 56, 0],
+        [10, 0, 2, 35, 0, 41, 56, 0],
+        [10, 0, 2, 35, 0, 41, 56, 0],
+    ],
+]
+
+
+@pytest.mark.parametrize("features", int_features)
+@pytest.mark.parametrize("pad_with", [None, 0.0, 0.1])
+@pytest.mark.parametrize("dtype", ["int8", "uint8", "int32", "float32", "int64", "float64"])
+class TestInterfaceDtypes:
+    """Unit tests to verify that AmplitudeEmbedding works correctly for features with all dtypes"""
+
+    @pytest.mark.autograd
+    def test_autograd(self, tol, features, pad_with, dtype):
+        """Tests autograd tensors."""
+
+        dtype = getattr(pnp, dtype)
+        features = pnp.array(features, dtype=dtype)
+
+        dev = qml.device("default.qubit")
+
+        circuit = qml.QNode(circuit_template, dev, interface="autograd")
+        circuit2 = qml.QNode(circuit_decomposed, dev)
+
+        res = circuit(features, pad_with, normalize=True)
+        res2 = circuit2(features, pad_with, normalize=True)
+
+        assert qml.math.allclose(res, res2, atol=tol, rtol=0)
+
+    @pytest.mark.jax
+    def test_jax(self, tol, features, pad_with, dtype):
+        """Tests jax tensors."""
+        import jax.numpy as jnp
+
+        dtype = getattr(jnp, dtype)
+        features = jnp.array(features, dtype=dtype)
+
+        dev = qml.device("default.qubit")
+
+        circuit = qml.QNode(circuit_template, dev, interface="jax")
+        circuit2 = qml.QNode(circuit_decomposed, dev)
+
+        res = circuit(features, pad_with, normalize=True)
+        res2 = circuit2(features, pad_with, normalize=True)
+
+        assert qml.math.allclose(res, res2, atol=tol, rtol=0)
+
+    @pytest.mark.jax
+    def test_jax_jit(self, tol, features, pad_with, dtype):
+        """Tests jax tensors with JIT compilation."""
+        import jax
+        import jax.numpy as jnp
+
+        dtype = getattr(jnp, dtype)
+        features = jnp.array(features, dtype=dtype)
+
+        dev = qml.device("default.qubit")
+
+        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, 2])
+
+        res = circuit(features, pad_with, normalize=True)
+        res2 = circuit2(features, pad_with, normalize=True)
+
+        assert qml.math.allclose(res, res2, atol=tol, rtol=0)
+
+    @pytest.mark.tf
+    def test_tf(self, tol, features, pad_with, dtype):
+        """Tests tensorflow tensors."""
+        import tensorflow as tf
+
+        dtype = getattr(tf, dtype)
+        features = tf.Variable(features, dtype=dtype)
+
+        dev = qml.device("default.qubit")
+
+        circuit = qml.QNode(circuit_template, dev, interface="tensorflow")
+        circuit2 = qml.QNode(circuit_decomposed, dev)
+
+        res = circuit(features, pad_with, normalize=True)
+        res2 = circuit2(features, pad_with, normalize=True)
+
+        assert qml.math.allclose(res, res2, atol=tol, rtol=0)
+
+    @pytest.mark.tf
+    def test_tf_jit(self, tol, features, pad_with, dtype):
+        """Tests tensorflow tensors with JIT compilation."""
+        import tensorflow as tf
+
+        dtype = getattr(tf, dtype)
+        features = tf.Variable(features, dtype=dtype)
+
+        dev = qml.device("default.qubit")
+
+        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, pad_with, normalize=True)
+        res2 = circuit2(features, pad_with, normalize=True)
+
+        assert qml.math.allclose(res, res2, atol=tol, rtol=0)
+
+    @pytest.mark.torch
+    def test_torch(self, tol, features, pad_with, dtype):
+        """Tests Torch tensors."""
+        import torch
+
+        dtype = getattr(torch, dtype)
+        features = torch.tensor(features, dtype=dtype)
+
+        dev = qml.device("default.qubit")
+
+        circuit = qml.QNode(circuit_template, dev, interface="torch")
+        circuit2 = qml.QNode(circuit_decomposed, dev)
+
+        res = circuit(features, pad_with, normalize=True)
+        res2 = circuit2(features, pad_with, normalize=True)
+
+        assert qml.math.allclose(res, res2, atol=tol, rtol=0)