Repository: PennyLaneAI/pennylane
Pull request #3298 (2022-11-09 03:24:16)

Daily rc sync to master 2022-11-09
Automatic sync from the release candidate to master during a feature freeze.

codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3298?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#3298](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3298?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (8cc83ac) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/757dcdc664145adeea0ee4f4ba15305ebdbdc077?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (757dcdc) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

```diff
@@           Coverage Diff           @@
##           master    #3298   +/-   ##
=======================================
  Coverage   99.70%   99.70%           
=======================================
  Files         299      299           
  Lines       26114    26121    +7     
=======================================
+ Hits        26037    26044    +7     
  Misses         77       77           
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3298?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/interfaces/jax\_jit.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3298/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvamF4X2ppdC5weQ==) | `100.00% <ø> (ø)` | |
| [pennylane/math/single\_dispatch.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3298/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvc2luZ2xlX2Rpc3BhdGNoLnB5) | `99.13% <100.00%> (+<0.01%)` | :arrow_up: |
| [pennylane/ops/qubit/observables.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3298/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9xdWJpdC9vYnNlcnZhYmxlcy5weQ==) | `100.00% <100.00%> (ø)` | |
| [pennylane/pauli/pauli\_arithmetic.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3298/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3BhdWxpL3BhdWxpX2FyaXRobWV0aWMucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/templates/state\_preparations/mottonen.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3298/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RlbXBsYXRlcy9zdGF0ZV9wcmVwYXJhdGlvbnMvbW90dG9uZW4ucHk=) | `100.00% <100.00%> (ø)` | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)


Patch:
diff --git a/doc/releases/changelog-0.27.0.md b/doc/releases/changelog-0.27.0.md
index 92bbb0226..d537a35e1 100644
--- a/doc/releases/changelog-0.27.0.md
+++ b/doc/releases/changelog-0.27.0.md
@@ -147,6 +147,11 @@
   [(#3261)](https://github.com/PennyLaneAI/pennylane/pull/3261)
 
   ```python
+  import jax
+  from jax import numpy as jnp
+  from jax.config import config
+  config.update("jax_enable_x64", True)
+
   dev = qml.device("lightning.qubit", wires=2)
 
   @jax.jit
@@ -157,17 +162,17 @@
       qml.CNOT(wires=[0, 1])
       return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))
 
-  x = jnp.array(1.0, dtype=jnp.float32)
-  y = jnp.array(2.0, dtype=jnp.float32)
+  x = jnp.array(1.0)
+  y = jnp.array(2.0)
   ```
 
   ```pycon
   >>> jax.jacobian(circuit, argnums=[0, 1])(x, y)
-  (DeviceArray([-0.84147096,  0.3501755 ], dtype=float32),
-   DeviceArray([ 4.474455e-18, -4.912955e-01], dtype=float32))
+  (DeviceArray([-0.84147098,  0.35017549], dtype=float64, weak_type=True),
+   DeviceArray([ 4.47445479e-18, -4.91295496e-01], dtype=float64, weak_type=True))
   ```
 
-  Note that this change depends on `jax.pure_callback`, which requires `jax==0.3.17`.
+  Note that this change depends on `jax.pure_callback`, which requires `jax>=0.3.17`.
 
 * The `QNode` class now accepts an `auto` interface, which automatically detects the interface of the given input.
   [(#3132)](https://github.com/PennyLaneAI/pennylane/pull/3132)
@@ -515,7 +520,7 @@
 
 * Extended `qml.equal` function to MeasurementProcesses
   [(#3189)](https://github.com/PennyLaneAI/pennylane/pull/3189)
-  
+
 * `_multi_dispatch` functionality has been moved inside the `get_interface` function. This function
   can now be called with one or multiple tensors as arguments.
   [(#3136)](https://github.com/PennyLaneAI/pennylane/pull/3136)
@@ -547,10 +552,11 @@
   >>> qml.math.get_interface(torch_scalar, torch_tensor, numpy_tensor)
   'torch'
   ```
-  
-* `qml.drawer.draw.draw_mpl` now accepts a `style` kwarg to select a style for plotting, rather than calling 
-  `qml.drawer.use_style(style)` before plotting. Setting a style for `draw_mpl` does not change the global 
-  configuration for matplotlib plotting. If no `style` is passed, the function defaults 
+
+* `qml.drawer.draw.draw_mpl` now accepts a `style` kwarg to select a style for plotting, rather than calling
+  `qml.drawer.use_style(style)` before plotting. Setting a style for `draw_mpl` does not change the global
+  configuration for matplotlib plotting. If no `style` is passed, the function defaults
+
   to plotting with the `black_white` style.
   [(#3247)](https://github.com/PennyLaneAI/pennylane/pull/3247)
 
@@ -591,6 +597,9 @@
 
 <h3>Bug fixes</h3>
 
+* `qml.SparseHamiltonian` now validates the size of the input matrix.
+  [(#3278)](https://github.com/PennyLaneAI/pennylane/pull/3278)
+
 * Fixed a bug where `qml.sample()` and `qml.counts()` would output incorrect results when mixed with measurements whose
   operators do not qubit-wise commute with computational basis projectors.
   [(#3207)](https://github.com/PennyLaneAI/pennylane/pull/3207)
@@ -635,6 +644,9 @@
 * Fixed a bug where `qml.BasisEmbedding` was not jit-compilable with JAX.
   [(#3239)](https://github.com/PennyLaneAI/pennylane/pull/3239)
 
+* Fixed a bug where `qml.MottonenStatePreparation` was not jit-compilable with JAX.
+  [(#3260)](https://github.com/PennyLaneAI/pennylane/pull/3260)
+
 * Fixed a bug where `qml.expval(qml.Hamiltonian())` would not raise an error
   if the Hamiltonian involved some wires that are not present on the device.
   [(#3266)](https://github.com/PennyLaneAI/pennylane/pull/3266)
diff --git a/pennylane/interfaces/jax_jit.py b/pennylane/interfaces/jax_jit.py
index e6684ec0c..768c574d1 100644
--- a/pennylane/interfaces/jax_jit.py
+++ b/pennylane/interfaces/jax_jit.py
@@ -31,11 +31,11 @@ dtype = jnp.float64
 
 def _validate_jax_version():
     if semantic_version.match("<0.3.17", jax.__version__) or semantic_version.match(
-        "<0.3.17", jax.lib.__version__
+        "<0.3.15", jax.lib.__version__
     ):
         msg = (
-            "The JAX JIT interface of PennyLane requires version 0.3.17 or higher for JAX and JAX lib. "
-            "Please upgrade these packages."
+            "The JAX JIT interface of PennyLane requires version 0.3.17 or higher for JAX "
+            "and 0.3.15 or higher JAX lib. Please upgrade these packages."
         )
         raise InterfaceUnsupportedError(msg)
 
diff --git a/pennylane/math/single_dispatch.py b/pennylane/math/single_dispatch.py
index 2202d4a17..2d43a7c16 100644
--- a/pennylane/math/single_dispatch.py
+++ b/pennylane/math/single_dispatch.py
@@ -201,6 +201,7 @@ ar.autoray._SUBMODULE_ALIASES["tensorflow", "moveaxis"] = "tensorflow.experiment
 ar.autoray._SUBMODULE_ALIASES["tensorflow", "sinc"] = "tensorflow.experimental.numpy"
 ar.autoray._SUBMODULE_ALIASES["tensorflow", "isclose"] = "tensorflow.experimental.numpy"
 ar.autoray._SUBMODULE_ALIASES["tensorflow", "atleast_1d"] = "tensorflow.experimental.numpy"
+ar.autoray._SUBMODULE_ALIASES["tensorflow", "all"] = "tensorflow.experimental.numpy"
 
 ar.autoray._FUNC_ALIASES["tensorflow", "arcsin"] = "asin"
 ar.autoray._FUNC_ALIASES["tensorflow", "arccos"] = "acos"
diff --git a/pennylane/ops/qubit/observables.py b/pennylane/ops/qubit/observables.py
index 7d1e50ca8..5138d7ec3 100644
--- a/pennylane/ops/qubit/observables.py
+++ b/pennylane/ops/qubit/observables.py
@@ -252,6 +252,11 @@ class SparseHamiltonian(Observable):
         if not isinstance(H, csr_matrix):
             raise TypeError("Observable must be a scipy sparse csr_matrix.")
         super().__init__(H, wires=wires, do_queue=do_queue, id=id)
+        mat_len = 2 ** len(self.wires)
+        if H.shape != (mat_len, mat_len):
+            raise ValueError(
+                f"Sparse Matrix must be of shape ({mat_len}, {mat_len}). Got {H.shape}."
+            )
 
     def label(self, decimals=None, base_label=None, cache=None):
         return super().label(decimals=decimals, base_label=base_label or "𝓗", cache=cache)
diff --git a/pennylane/pauli/pauli_arithmetic.py b/pennylane/pauli/pauli_arithmetic.py
index 1927df00f..b2001661d 100644
--- a/pennylane/pauli/pauli_arithmetic.py
+++ b/pennylane/pauli/pauli_arithmetic.py
@@ -143,7 +143,7 @@ class PauliWord(dict):
     def __str__(self):
         """String representation of a PauliWord."""
         if len(self) == 0:
-            return "()"
+            return "I"
         return " @ ".join(f"{op}({w})" for w, op in self.items())
 
     def __repr__(self):
@@ -224,6 +224,8 @@ class PauliSentence(dict):
 
     def __str__(self):
         """String representation of the PauliSentence."""
+        if len(self) == 0:
+            return "I"
         return "\n+ ".join(f"{coeff} * {str(pw)}" for pw, coeff in self.items())
 
     def __repr__(self):
diff --git a/pennylane/templates/state_preparations/mottonen.py b/pennylane/templates/state_preparations/mottonen.py
index 3cd71d536..602641c57 100644
--- a/pennylane/templates/state_preparations/mottonen.py
+++ b/pennylane/templates/state_preparations/mottonen.py
@@ -88,7 +88,7 @@ def compute_theta(alpha):
         for j in range(len(M_trans[0])):
             M_trans[i, j] = _matrix_M_entry(j, i)
 
-    theta = qml.math.dot(M_trans, alpha.T).T
+    theta = qml.math.transpose(qml.math.dot(M_trans, qml.math.transpose(alpha)))
 
     return theta / 2**k
 
@@ -123,7 +123,7 @@ def _apply_uniform_rotation_dagger(gate, alpha, control_wires, target_wire):
     gray_code_rank = len(control_wires)
 
     if gray_code_rank == 0:
-        if qml.math.all(theta[..., 0] != 0.0):
+        if qml.math.is_abstract(theta) or qml.math.all(theta[..., 0] != 0.0):
             op_list.append(gate(theta[..., 0], wires=[target_wire]))
         return op_list
 
@@ -136,7 +136,7 @@ def _apply_uniform_rotation_dagger(gate, alpha, control_wires, target_wire):
     ]
 
     for i, control_index in enumerate(control_indices):
-        if qml.math.all(theta[..., i] != 0.0):
+        if qml.math.is_abstract(theta) or qml.math.all(theta[..., i] != 0.0):
             op_list.append(gate(theta[..., i], wires=[target_wire]))
         op_list.append(qml.CNOT(wires=[control_wires[control_index], target_wire]))
     return op_list
@@ -309,11 +309,12 @@ class MottonenStatePreparation(Operation):
                     f"State vectors must be of length {2 ** len(wires)} or less; vector {i} has length {n_amplitudes}."
                 )
 
-            norm = qml.math.sum(qml.math.abs(state) ** 2)
-            if not qml.math.allclose(norm, 1.0, atol=1e-3):
-                raise ValueError(
-                    f"State vectors have to be of norm 1.0, vector {i} has norm {norm}"
-                )
+            if not qml.math.is_abstract(state):
+                norm = qml.math.sum(qml.math.abs(state) ** 2)
+                if not qml.math.allclose(norm, 1.0, atol=1e-3):
+                    raise ValueError(
+                        f"State vectors have to be of norm 1.0, vector {i} has norm {norm}"
+                    )
 
         super().__init__(state_vector, wires=wires, do_queue=do_queue, id=id)
 
@@ -364,7 +365,7 @@ class MottonenStatePreparation(Operation):
             op_list.extend(_apply_uniform_rotation_dagger(qml.RY, alpha_y_k, control, target))
 
         # If necessary, apply inverse z rotation cascade to prepare correct phases of amplitudes
-        if not qml.math.allclose(omega, 0):
+        if qml.math.is_abstract(omega) or not qml.math.allclose(omega, 0):
             for k in range(len(wires_reverse), 0, -1):
                 alpha_z_k = _get_alpha_z(omega, len(wires_reverse), k)
                 control = wires_reverse[k:]


Test patch:
diff --git a/tests/interfaces/test_jax.py b/tests/interfaces/test_jax.py
index 870b2c76b..adae2401f 100644
--- a/tests/interfaces/test_jax.py
+++ b/tests/interfaces/test_jax.py
@@ -33,9 +33,16 @@ from pennylane.interfaces.jax_jit import _execute_with_fwd
 
 
 @pytest.mark.parametrize(
-    "version, should_raise", [("0.3.16", True), ("0.3.17", False), ("0.3.18", False)]
+    "version, package, should_raise",
+    [
+        ("0.3.16", jax, True),
+        ("0.3.17", jax, False),
+        ("0.3.18", jax, False),
+        ("0.3.14", jax.lib, True),
+        ("0.3.15", jax.lib, False),
+        ("0.3.16", jax.lib, False),
+    ],
 )
-@pytest.mark.parametrize("package", [jax, jax.lib])
 def test_raise_version_error(package, version, should_raise, monkeypatch):
     """Test JAX version error"""
     a = jnp.array([0.1, 0.2])
@@ -49,7 +56,7 @@ def test_raise_version_error(package, version, should_raise, monkeypatch):
         m.setattr(package, "__version__", version)
 
         if should_raise:
-            msg = "requires version 0.3.17 or higher for JAX and JAX lib"
+            msg = "requires version 0.3.17 or higher for JAX and 0.3.15 or higher JAX lib"
             with pytest.raises(InterfaceUnsupportedError, match=msg):
                 execute([tape], dev, gradient_fn=param_shift, interface="jax-jit")
         else:
diff --git a/tests/ops/functions/test_eigvals.py b/tests/ops/functions/test_eigvals.py
index a0f4a9074..732b38e6c 100644
--- a/tests/ops/functions/test_eigvals.py
+++ b/tests/ops/functions/test_eigvals.py
@@ -203,23 +203,24 @@ class TestSingleOperation:
     )
     def test_sparse_hamiltonian(self, row, col, dat, val_ref):
         """Test that the eigenvalues of a sparse Hamiltonian are correctly returned"""
-        # N x N matrix with N = 20
-        h_sparse = scipy.sparse.csr_matrix((dat, (row, col)), shape=(len(row), len(col)))
-        h_sparse = qml.SparseHamiltonian(h_sparse, wires=all)
+        # N x N matrix with N = 16
+        h_mat = scipy.sparse.csr_matrix((dat, (row, col)), shape=(16, 16))
+        h_sparse = qml.SparseHamiltonian(h_mat, wires=range(4))
+
+        dense_mat = h_mat.todense()
+        dense_eigvals = np.sort(np.linalg.eigvals(dense_mat))
 
         # k = 1  (< N-1) scipy.sparse.linalg is used:
         val_groundstate = qml.eigvals(h_sparse, k=1)
-        # k = 18  (< N-1) scipy.sparse.linalg is used:
-        val_n_sparse = qml.eigvals(h_sparse, k=len(val_ref) - 2)
-        # k = 20 (> N-1) qml.math.linalg is used:
-        val_all = qml.eigvals(h_sparse, k=len(val_ref))
-        # k = 19 (= N-1) qml.math.linalg is used:
-        val_n_dense = qml.eigvals(h_sparse, k=len(val_ref) - 1)
-
-        assert np.allclose(val_groundstate, val_ref[0])
-        assert np.allclose(np.sort(val_n_sparse), val_ref[0:18])
-        assert np.allclose(np.sort(val_all), val_ref)
-        assert np.allclose(np.sort(val_n_dense), val_ref)
+        assert np.allclose(val_groundstate, dense_eigvals[0])
+
+        # k = 14  (< N-1) scipy.sparse.linalg is used:
+        val_n_sparse = np.sort(qml.eigvals(h_sparse, k=14))
+        assert np.allclose(val_n_sparse, dense_eigvals[0:-2])
+
+        # k = 16 (> N-1) qml.math.linalg is used:
+        val_all = np.sort(qml.eigvals(h_sparse, k=16))
+        assert np.allclose(val_all, dense_eigvals)
 
 
 class TestMultipleOperations:
diff --git a/tests/ops/qubit/test_sparse.py b/tests/ops/qubit/test_sparse.py
index 8ee3eb514..da13e17b9 100644
--- a/tests/ops/qubit/test_sparse.py
+++ b/tests/ops/qubit/test_sparse.py
@@ -62,6 +62,14 @@ class TestSparse:
         H = qml.SparseHamiltonian(csr_matrix(np.array([[1, 0], [-1.5, 0]])), 1)
         assert H.label() == "𝓗"
 
+    def test_valueeror_shape(self):
+        """Test that iniatilization raises a ValueError if the shape is not correct for the number of wires."""
+        mat = csr_matrix([[1, 0], [0, 1]])
+        with pytest.raises(
+            ValueError, match=r"Sparse Matrix must be of shape \(4, 4\). Got \(2, 2\)."
+        ):
+            qml.SparseHamiltonian(mat, wires=(0, 1))
+
     @pytest.mark.parametrize("sparse_hamiltonian", SPARSE_HAMILTONIAN_TEST_DATA)
     def test_sparse_typeerror(self, sparse_hamiltonian):
         """Test that the SparseHamiltonian class raises a TypeError on incorrect inputs."""
@@ -69,12 +77,12 @@ class TestSparse:
         sparse_hamiltonian = coo_matrix(mat)
 
         with pytest.raises(TypeError, match="Observable must be a scipy sparse csr_matrix"):
-            qml.SparseHamiltonian(sparse_hamiltonian)
+            qml.SparseHamiltonian(sparse_hamiltonian, wires=0)
 
     @pytest.mark.parametrize("sparse_hamiltonian", SPARSE_HAMILTONIAN_TEST_DATA)
     def test_sparse_matrix(self, sparse_hamiltonian, tol):
         """Test that the matrix property of the SparseHamiltonian class returns the correct matrix."""
-        num_wires = len(sparse_hamiltonian[0])
+        num_wires = int(np.log2(len(sparse_hamiltonian[0])))
         sparse_hamiltonian_csr = csr_matrix(sparse_hamiltonian)
         res_dynamic = qml.SparseHamiltonian(
             sparse_hamiltonian_csr, range(num_wires)
@@ -88,7 +96,7 @@ class TestSparse:
     @pytest.mark.parametrize("sparse_hamiltonian", SPARSE_HAMILTONIAN_TEST_DATA)
     def test_matrix(self, sparse_hamiltonian, tol):
         """Test that the matrix property of the SparseHamiltonian class returns the correct matrix."""
-        num_wires = len(sparse_hamiltonian[0])
+        num_wires = int(np.log2(len(sparse_hamiltonian[0])))
         sparse_hamiltonian_csr = csr_matrix(sparse_hamiltonian)
         res_dynamic = qml.SparseHamiltonian(sparse_hamiltonian_csr, range(num_wires)).matrix()
         res_static = qml.SparseHamiltonian.compute_matrix(sparse_hamiltonian_csr)
diff --git a/tests/pauli/test_pauli_arithmetic.py b/tests/pauli/test_pauli_arithmetic.py
index 34a43209d..ec8ca8795 100644
--- a/tests/pauli/test_pauli_arithmetic.py
+++ b/tests/pauli/test_pauli_arithmetic.py
@@ -96,7 +96,7 @@ class TestPauliWord:
         (pw1, "X(1) @ Y(2)"),
         (pw2, "X(a) @ X(b) @ Z(c)"),
         (pw3, "Z(0) @ Z(b) @ Z(c)"),
-        (pw4, "()"),
+        (pw4, "I"),
     )
 
     @pytest.mark.parametrize("pw, str_rep", tup_pw_str)
@@ -185,8 +185,9 @@ class TestPauliSentence:
             ps2,
             "-1.23 * X(1) @ Y(2)\n" "+ (-0-4j) * X(a) @ X(b) @ Z(c)\n" "+ 0.5 * Z(0) @ Z(b) @ Z(c)",
         ),
-        (ps3, "-0.5 * Z(0) @ Z(b) @ Z(c)\n" "+ 1 * ()"),
-        (ps4, "1 * ()"),
+        (ps3, "-0.5 * Z(0) @ Z(b) @ Z(c)\n" "+ 1 * I"),
+        (ps4, "1 * I"),
+        (ps5, "I"),
     )
 
     @pytest.mark.parametrize("ps, str_rep", tup_ps_str)
diff --git a/tests/templates/test_state_preparations/test_mottonen_state_prep.py b/tests/templates/test_state_preparations/test_mottonen_state_prep.py
index f93379b53..d1e8173a2 100644
--- a/tests/templates/test_state_preparations/test_mottonen_state_prep.py
+++ b/tests/templates/test_state_preparations/test_mottonen_state_prep.py
@@ -357,22 +357,76 @@ class TestGradient:
         qml.grad(circuit)(state_vector)
 
 
+@pytest.mark.parametrize(
+    "inputs, expected",
+    [
+        (
+            [0.0, 0.7, 0.7, 0.0],
+            [0.0, 0.5, 0.5, 0.0],
+        ),
+        ([0.1, 0.0, 0.0, 0.1], [0.5, 0.0, 0.0, 0.5]),
+    ],
+)
 class TestCasting:
     """Test that the Mottonen state preparation ensures the compatibility with
     interfaces by using casting'"""
 
+    @pytest.mark.jax
+    def test_jax(self, inputs, expected):
+        """Test that MottonenStatePreparation can be correctly used with the JAX interface."""
+        import jax
+        from jax import numpy as jnp
+
+        inputs = jnp.array(inputs)
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, interface="jax")
+        def circuit(inputs):
+            qml.MottonenStatePreparation(inputs, wires=[0, 1])
+            return qml.probs(wires=[0, 1])
+
+        inputs = inputs / jnp.linalg.norm(inputs)
+        res = circuit(inputs)
+        assert np.allclose(res, expected, atol=1e-6, rtol=0)
+
+    @pytest.mark.jax
+    def test_jax_jit(self, inputs, expected):
+        """Test that MottonenStatePreparation can be correctly used with the JAX-JIT interface."""
+        import jax
+        from jax import numpy as jnp
+
+        inputs = jnp.array(inputs)
+        dev = qml.device("default.qubit", wires=2)
+
+        @jax.jit
+        @qml.qnode(dev, interface="jax")
+        def circuit(inputs):
+            qml.MottonenStatePreparation(inputs, wires=[0, 1])
+            return qml.probs(wires=[0, 1])
+
+        inputs = inputs / jnp.linalg.norm(inputs)
+        res = circuit(inputs)
+        assert np.allclose(res, expected, atol=1e-6, rtol=0)
+
+    @pytest.mark.tf
+    def test_tensorflow(self, inputs, expected):
+        """Test that MottonenStatePreparation can be correctly used with the TensorFlow interface."""
+        import tensorflow as tf
+
+        inputs = tf.Variable(inputs)
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, interface="tf")
+        def circuit(inputs):
+            qml.MottonenStatePreparation(inputs, wires=[0, 1])
+            return qml.probs(wires=[0, 1])
+
+        inputs = inputs / tf.linalg.norm(inputs)
+        res = circuit(inputs)
+        assert np.allclose(res, expected, atol=1e-6, rtol=0)
+
     @pytest.mark.torch
-    @pytest.mark.parametrize(
-        "inputs, expected",
-        [
-            (
-                [0.0, 0.7, 0.7, 0.0],
-                [0.0, 0.5, 0.5, 0.0],
-            ),
-            ([0.1, 0.0, 0.0, 0.1], [0.5, 0.0, 0.0, 0.5]),
-        ],
-    )
-    def test_scalar_torch(self, inputs, expected):
+    def test_torch(self, inputs, expected):
         """Test that MottonenStatePreparation can be correctly used with the Torch interface."""
         import torch