Repository: PennyLaneAI/pennylane
Pull request #6896 (2025-01-28 22:27:33)

Adding decomposition MPSPrep
### Before submitting

Please complete the following checklist when submitting a PR:

- [x] All new features must include a unit test.
      If you've fixed a bug or added code that should be tested, add a test to the
      test directory!

- [x] All new functions and code must be clearly commented and documented.
      If you do make documentation changes, make sure that the docs build and
      render correctly by running `make docs`.

- [x] Ensure that the test suite passes, by running `make test`.

- [x] Add a new entry to the `doc/releases/changelog-dev.md` file, summarizing the
      change, and including a link back to the PR.

- [x] The PennyLane source code conforms to
      [PEP8 standards](https://www.python.org/dev/peps/pep-0008/).
      We check all of our code against [Pylint](https://www.pylint.org/).
      To lint modified files, simply `pip install pylint`, and then
      run `pylint pennylane/path/to/file.py`.

When all the above are checked, delete everything above the dashed
line and fill in the pull request template.

------------------------------------------------------------------------------------------------------------

**Context:**

[SC-83733]
[SC-81348]
MPSPrep only works in "lightning.tensor". With this decomposition, it will work in any device

**Description of the Change:**

**Benefits:**

**Possible Drawbacks:**

**Related GitHub Issues:**



** aux code **

useful code for review

```python
def mps_right_canonical(state, max_bond_dim):
    """
    Convert a state vector into a right-canonical MPS.

    Parameters:
        state (array): 1D state vector of dimension 2^L.
        max_bond_dim (int): Maximum bond dimension.

    Returns:
        list of array: MPS tensors in right-canonical form.
    """
    d = 2
    L = int(np.log2(len(state)))
    psi = state.reshape([d] * L)
    psi_rev = np.transpose(psi, axes=range(L)[::-1])

    mps_rev = []
    current = psi_rev.copy()
    left_dim = 1

    for _ in range(L - 1):
        current = current.reshape(left_dim * d, -1)
        U, S, Vh = np.linalg.svd(current, full_matrices=False)
        chi = min(len(S), max_bond_dim)
        U = U[:, :chi]
        S = S[:chi]
        Vh = Vh[:chi, :]
        mps_rev.append(U.reshape(left_dim, d, chi))
        left_dim = chi
        current = np.dot(np.diag(S), Vh)

    mps_rev.append(current.reshape(left_dim, d, 1))
    mps = [np.transpose(A, axes=(2, 1, 0)) for A in mps_rev[::-1]]
    mps[0] = mps[0][0]
    mps[-1] = mps[-1][:, :, 0]

    return mps
```




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.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6896?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 99.59%. Comparing base [(`06cd2f0`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/06cd2f0d01ff67495acc566ff0244202da7e48c8?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`67f04b6`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/67f04b64cc51caed9e6c15c4c3a0f2c958861df3?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff           @@
##           master    #6896   +/-   ##
=======================================
  Coverage   99.59%   99.59%           
=======================================
  Files         478      478           
  Lines       45306    45359   +53     
=======================================
+ Hits        45124    45177   +53     
  Misses        182      182           
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6896?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

KetpuntoG: I have deleted the controversial long test and added a smaller one

Patch:
diff --git a/doc/introduction/templates.rst b/doc/introduction/templates.rst
index b5b84b8b9..46073098c 100644
--- a/doc/introduction/templates.rst
+++ b/doc/introduction/templates.rst
@@ -114,6 +114,11 @@ state preparation is typically used as the first operation.
     :description: :doc:`CosineWindow <../code/api/pennylane.CosineWindow>`
     :figure: _static/templates/state_preparations/thumbnail_cosine_window.png
 
+.. gallery-item::
+    :description: :doc:`MPSPrep <../code/api/pennylane.MPSPrep>`
+    :figure: _static/templates/tensornetworks/MPS_template.png
+
+
 .. raw:: html
 
         <div style='clear:both'></div>
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index cbe5de5f9..aa3b8157e 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -235,6 +235,10 @@
 * The requested `diff_method` is now validated when program capture is enabled.
   [(#6852)](https://github.com/PennyLaneAI/pennylane/pull/6852)
 
+* The template `MPSPrep` now has a gate decomposition. This enables its use with any device.
+  The `right_canonicalize_mps` function has also been added to transform an MPS into its right-canonical form.
+  [(#6896)](https://github.com/PennyLaneAI/pennylane/pull/6896)
+
 * The `qml.clifford_t_decomposition` has been improved to use less gates when decomposing `qml.PhaseShift`.
   [(#6842)](https://github.com/PennyLaneAI/pennylane/pull/6842)
 
diff --git a/pennylane/templates/state_preparations/__init__.py b/pennylane/templates/state_preparations/__init__.py
index bc2f902c3..6adff7e6b 100644
--- a/pennylane/templates/state_preparations/__init__.py
+++ b/pennylane/templates/state_preparations/__init__.py
@@ -21,4 +21,4 @@ from .basis_qutrit import QutritBasisStatePreparation
 from .cosine_window import CosineWindow
 from .mottonen import MottonenStatePreparation
 from .superposition import Superposition
-from .state_prep_mps import MPSPrep
+from .state_prep_mps import MPSPrep, right_canonicalize_mps
diff --git a/pennylane/templates/state_preparations/state_prep_mps.py b/pennylane/templates/state_preparations/state_prep_mps.py
index 92fd7edda..1a3549eba 100644
--- a/pennylane/templates/state_preparations/state_prep_mps.py
+++ b/pennylane/templates/state_preparations/state_prep_mps.py
@@ -15,27 +15,243 @@
 Contains the MPSPrep template.
 """
 
+import numpy as np
+
 import pennylane as qml
 from pennylane.operation import Operation
 from pennylane.wires import Wires
 
 
+def _validate_mps_shape(mps):
+    r"""Validate that the MPS dimensions are correct.
+
+    Args:
+        mps (list[TensorLike]): List of tensors representing the MPS.
+    """
+
+    # Validate the shape and dimensions of the first tensor
+    assert qml.math.isclose(
+        len(qml.math.shape(mps[0])), 2
+    ), "The first tensor must have exactly 2 dimensions."
+    dj0, dj2 = qml.math.shape(mps[0])
+    assert qml.math.isclose(dj0, 2), "The first dimension of the first tensor must be exactly 2."
+    assert qml.math.log2(
+        dj2
+    ).is_integer(), "The second dimension of the first tensor must be a power of 2."
+
+    # Validate the shapes of the intermediate tensors
+    for i, array in enumerate(mps[1:-1], start=1):
+        shape = qml.math.shape(array)
+        assert qml.math.isclose(len(shape), 3), f"Tensor {i} must have exactly 3 dimensions."
+        new_dj0, new_dj1, new_dj2 = shape
+        assert qml.math.isclose(
+            new_dj1, 2
+        ), f"The second dimension of tensor {i} must be exactly 2."
+        assert qml.math.log2(
+            new_dj0
+        ).is_integer(), f"The first dimension of tensor {i} must be a power of 2."
+        assert qml.math.isclose(
+            new_dj1, 2
+        ), f"The second dimension of tensor {i} must be exactly 2."
+        assert qml.math.log2(
+            new_dj2
+        ).is_integer(), f"The third dimension of tensor {i} must be a power of 2."
+        assert qml.math.isclose(
+            new_dj0, dj2
+        ), f"Dimension mismatch: tensor {i}'s first dimension does not match the previous third dimension."
+        dj2 = new_dj2
+
+    # Validate the shape and dimensions of the last tensor
+    assert qml.math.isclose(
+        len(qml.math.shape(mps[-1])), 2
+    ), "The last tensor must have exactly 2 dimensions."
+    new_dj0, new_dj1 = qml.math.shape(mps[-1])
+    assert new_dj1 == 2, "The second dimension of the last tensor must be exactly 2."
+    assert qml.math.log2(
+        new_dj0
+    ).is_integer(), "The first dimension of the last tensor must be a power of 2."
+    assert qml.math.isclose(
+        new_dj0, dj2
+    ), "Dimension mismatch: the last tensor's first dimension does not match the previous third dimension."
+
+
+def right_canonicalize_mps(mps):
+    r"""Transform a matrix product state (MPS) into its right-canonical form.
+
+    A right-canonicalized MPS is a matrix product state in which the constituent tensors, :math:`A^{(j)}`, satisfy
+    the following orthonormality condition [Eq. (21) of `arXiv:2310.18410 <https://arxiv.org/pdf/2310.18410>`_]:
+
+    .. math::
+
+        \sum_{d_{j,1}, d_{j,2}} A^{(j)}_{d_{j, 0}, d_{j, 1}, d_{j, 2}} \left( A^{(j)}_{d'_{j, 0}, d_{j, 1}, d_{j, 2}} \right)^* = \delta_{d_{j, 0}, d'_{j, 0}},
+
+    where :math:`d_{i,j}` denotes the :math:`j` dimension of the :math:`i` tensor and :math:`\delta` is the Kronecker delta.
+
+
+    Args:
+        mps (list[TensorLike]): List of tensors representing the MPS.
+
+    Returns:
+        List of tensors representing the MPS in right-canonical form with the same dimensions as the initial MPS.
+
+    .. seealso:: :class:`~.MPSPrep`.
+
+    **Example**
+
+    .. code-block::
+
+        n_sites = 4
+
+        mps = ([np.ones((2, 4))] +
+               [np.ones((4, 2, 4)) for _ in range(1, n_sites - 1)] +
+               [np.ones((4, 2))])
+
+        mps_rc = qml.right_canonicalize_mps(mps)
+
+        # Check that the right-canonical definition is fulfilled
+        for i in range(1, n_sites - 1):
+            tensor = mps_rc[i]
+            contraction_matrix = np.tensordot(tensor, tensor.conj(), axes=([1, 2], [1, 2]))
+            assert np.allclose(contraction_matrix, np.eye(tensor.shape[0]))
+
+    .. details::
+        :title: Usage Details
+
+        The input MPS must be a list of :math:`n` tensors :math:`[A^{(0)}, ..., A^{(n-1)}]`
+        with shapes :math:`d_0, ..., d_{n-1}`, respectively. The first and last tensors have rank :math:`2`
+        while the intermediate tensors have rank :math:`3`.
+
+        The first tensor must have the shape :math:`d_0 = (d_{0,0}, d_{0,1})` where :math:`d_{0,0}`
+        and :math:`d_{0,1}`  correspond to the physical dimension of the site and an auxiliary bond
+        dimension connecting it to the next tensor, respectively.
+
+        The last tensor must have the shape :math:`d_{n-1} = (d_{n-1,0}, d_{n-1,1})` where :math:`d_{n-1,0}`
+        and :math:`d_{n-1,1}` represent the auxiliary dimension from the previous site and the physical
+        dimension of the site, respectively.
+
+        The intermediate tensors must have the shape :math:`d_j = (d_{j,0}, d_{j,1}, d_{j,2})`, where:
+
+        - :math:`d_{j,0}` is the bond dimension connecting to the previous tensor
+        - :math:`d_{j,1}` is the physical dimension for the site
+        - :math:`d_{j,2}` is the bond dimension connecting to the next tensor
+
+        Note that the bond dimensions must match between adjacent tensors such that :math:`d_{j-1,2} = d_{j,0}`.
+
+        Additionally, the physical dimension of the site should always be fixed at :math:`2`
+        (since the dimension of a qubit is :math:`2`), while the other dimensions must be powers of two.
+
+        The following example shows a valid MPS input containing four tensors with
+        dimensions :math:`[(2,2), (2,2,4), (4,2,2), (2,2)]` which satisfy the criteria described above.
+
+        .. code-block::
+
+            mps = [
+                np.array([[0.0, 0.107], [0.994, 0.0]]),
+                np.array(
+                    [
+                        [[0.0, 0.0, 0.0, -0.0], [1.0, 0.0, 0.0, -0.0]],
+                        [[0.0, 1.0, 0.0, -0.0], [0.0, 0.0, 0.0, -0.0]],
+                    ]
+                ),
+                np.array(
+                    [
+                        [[-1.0, 0.0], [0.0, 0.0]],
+                        [[0.0, 0.0], [0.0, 1.0]],
+                        [[0.0, -1.0], [0.0, 0.0]],
+                        [[0.0, 0.0], [1.0, 0.0]],
+                    ]
+                ),
+                np.array([[-1.0, -0.0], [-0.0, -1.0]]),
+            ]
+    """
+
+    _validate_mps_shape(mps)
+
+    mps = mps.copy()
+    mps[0] = mps[0].reshape((1, *mps[0].shape))
+    mps[-1] = mps[-1].reshape((*mps[-1].shape, 1))
+
+    if not qml.math.is_abstract(mps[0]):
+        is_right_canonical = True
+        for tensor in mps[1:-1]:
+            # Right-canonical definition
+            input_matrix = qml.math.tensordot(tensor, tensor.conj(), axes=([1, 2], [1, 2]))
+            if not qml.math.allclose(input_matrix, qml.math.eye(tensor.shape[0])):
+                is_right_canonical = False
+                break
+
+        if is_right_canonical:
+            mps[0] = mps[0][0]
+            mps[-1] = mps[-1][:, :, 0]
+            return mps
+
+    d_shapes = []
+    for tensor in mps[1:-1]:
+        d_shapes += tensor.shape
+
+    max_bond_dim = qml.math.max(d_shapes)
+
+    n_sites = len(mps)
+    output_mps = [None] * n_sites
+
+    # Procedure analogous to the left-canonical conversion but starting from the right and storing the Vd,
+    # where Vd is the right matrix in the Singular Value Decomposition (SVD)
+    for i in range(n_sites - 1, 0, -1):
+        chi_left, d, chi_right = mps[i].shape
+        input_matrix = mps[i].reshape(chi_left, d * chi_right)
+
+        u_matrix, s_diag, vd_matrix = qml.math.linalg.svd(input_matrix, full_matrices=False)
+
+        # Truncate SVD components if needed
+        chi_new = min(int(max_bond_dim), len(s_diag))
+        u_matrix = u_matrix[:, :chi_new]
+        s_diag = s_diag[:chi_new]
+        vd_matrix = vd_matrix[:chi_new, :]
+
+        # Store Vd reshaped as an MPS tensor in the output MPS
+        output_mps[i] = vd_matrix.reshape(chi_new, d, chi_right)
+
+        # Contract U with diag(S) and merge it with the preceding MPS tensor, preserving the canonical structure
+        mps[i - 1] = qml.math.tensordot(
+            mps[i - 1], u_matrix @ qml.math.diag(s_diag), axes=([2], [0])
+        )
+
+    output_mps[0] = mps[0][0]
+    output_mps[-1] = output_mps[-1][:, :, 0]
+
+    return output_mps
+
+
 class MPSPrep(Operation):
     r"""Prepares an initial state from a matrix product state (MPS) representation.
 
     .. note::
 
-        Currently, this operator can only be used with ``qml.device(“lightning.tensor”)``.
-
+        This operator is natively supported on the ``lightning.tensor`` device, which is designed to run MPS
+        structures efficiently. For other devices, this operation prepares the state vector represented by the
+        MPS using a gate-based decomposition from Eq. (23) in `arXiv:2310.18410
+        <https://arxiv.org/pdf/2310.18410>`_, which requires the right canonicalization of the MPS using
+        the :func:`~.right_canonicalize_mps` function and defining auxiliary qubits with ``work_wires``.
 
     Args:
-        mps (List[Array]):  list of arrays of rank-3 and rank-2 tensors representing an MPS state as a
-            product of site matrices. See the usage details section for more information.
+        mps (list[TensorLike]):  list of arrays of rank-3 and rank-2 tensors representing an MPS state
+            as a product of site matrices. See the usage details section for more information.
+
+        wires (Sequence[int]): wires that the template acts on. It should match the number of MPS tensors.
+        work_wires (Sequence[int]): list of extra qubits needed in the decomposition. If the maximum dimension
+            of the MPS tensors is :math:`2^k`, then :math:`k` ``work_wires`` will be needed. If no ``work_wires`` are given,
+            this operator can only be executed on the ``lightning.tensor`` device. Default is ``None``.
+
+        right_canonicalize (bool): indicates whether a conversion to right-canonical form should be performed to the MPS.
+            Default is ``False``.
 
-        wires (Sequence[int]): wires that the template acts on
+
+    .. seealso:: :func:`~.right_canonicalize_mps`.
 
     **Example**
 
+    Example using the ``lightning.tensor`` device:
+
     .. code-block::
 
         mps = [
@@ -61,6 +277,22 @@ class MPSPrep(Operation):
         [ 0.        +0.j -0.10705513+0.j  0.        +0.j  0.        +0.j
         0.        +0.j  0.        +0.j -0.99451217+0.j  0.        +0.j]
 
+    Example using the ``default.qubit`` device:
+
+    .. code-block::
+
+        dev = qml.device("default.qubit", wires=4)
+        @qml.qnode(dev)
+        def circuit():
+            qml.MPSPrep(mps, wires = [1,2,3], work_wires = [0])
+            return qml.state()
+
+    .. code-block:: pycon
+
+        >>> print(circuit()[:8])
+        [ 0.        +0.j -0.10705513+0.j  0.        +0.j  0.        +0.j
+        0.        +0.j  0.        +0.j -0.99451217+0.j  0.        +0.j]
+
     .. details::
         :title: Usage Details
 
@@ -112,55 +344,23 @@ class MPSPrep(Operation):
             ]
     """
 
-    def __init__(self, mps, wires, id=None):
+    def __init__(
+        self, mps, wires, work_wires=None, right_canonicalize=False, id=None
+    ):  # pylint: disable=too-many-arguments,too-many-positional-arguments
 
-        # Validate the shape and dimensions of the first tensor
-        assert qml.math.isclose(
-            len(qml.math.shape(mps[0])), 2
-        ), "The first tensor must have exactly 2 dimensions."
-        dj0, dj2 = qml.math.shape(mps[0])
-        assert qml.math.isclose(
-            dj0, 2
-        ), "The first dimension of the first tensor must be exactly 2."
-        assert qml.math.log2(
-            dj2
-        ).is_integer(), "The second dimension of the first tensor must be a power of 2."
-
-        # Validate the shapes of the intermediate tensors
-        for i, array in enumerate(mps[1:-1], start=1):
-            shape = qml.math.shape(array)
-            assert qml.math.isclose(len(shape), 3), f"Tensor {i} must have exactly 3 dimensions."
-            new_dj0, new_dj1, new_dj2 = shape
-            assert qml.math.isclose(
-                new_dj1, 2
-            ), f"The second dimension of tensor {i} must be exactly 2."
-            assert qml.math.log2(
-                new_dj0
-            ).is_integer(), f"The first dimension of tensor {i} must be a power of 2."
-            assert qml.math.isclose(
-                new_dj1, 2
-            ), f"The second dimension of tensor {i} must be exactly 2."
-            assert qml.math.log2(
-                new_dj2
-            ).is_integer(), f"The third dimension of tensor {i} must be a power of 2."
-            assert qml.math.isclose(
-                new_dj0, dj2
-            ), f"Dimension mismatch: tensor {i}'s first dimension does not match the previous third dimension."
-            dj2 = new_dj2
-
-        # Validate the shape and dimensions of the last tensor
-        assert qml.math.isclose(
-            len(qml.math.shape(mps[-1])), 2
-        ), "The last tensor must have exactly 2 dimensions."
-        new_dj0, new_dj1 = qml.math.shape(mps[-1])
-        assert new_dj1 == 2, "The second dimension of the last tensor must be exactly 2."
-        assert qml.math.log2(
-            new_dj0
-        ).is_integer(), "The first dimension of the last tensor must be a power of 2."
-        assert qml.math.isclose(
-            new_dj0, dj2
-        ), "Dimension mismatch: the last tensor's first dimension does not match the previous third dimension."
-        super().__init__(*mps, wires=wires, id=id)
+        _validate_mps_shape(mps)
+
+        self.hyperparameters["input_wires"] = qml.wires.Wires(wires)
+        self.hyperparameters["right_canonicalize"] = right_canonicalize
+
+        if work_wires:
+            self.hyperparameters["work_wires"] = qml.wires.Wires(work_wires)
+            all_wires = self.hyperparameters["input_wires"] + self.hyperparameters["work_wires"]
+        else:
+            self.hyperparameters["work_wires"] = None
+            all_wires = self.hyperparameters["input_wires"]
+
+        super().__init__(*mps, wires=all_wires, id=id)
 
     @property
     def mps(self):
@@ -168,7 +368,11 @@ class MPSPrep(Operation):
         return self.data
 
     def _flatten(self):
-        hyperparameters = (("wires", self.wires),)
+        hyperparameters = (
+            ("wires", self.hyperparameters["input_wires"]),
+            ("work_wires", self.hyperparameters["work_wires"]),
+            ("right_canonicalize", self.hyperparameters["right_canonicalize"]),
+        )
         return self.mps, hyperparameters
 
     @classmethod
@@ -177,8 +381,16 @@ class MPSPrep(Operation):
         return cls(data, **hyperparams_dict)
 
     def map_wires(self, wire_map):
-        new_wires = Wires([wire_map.get(wire, wire) for wire in self.wires])
-        return MPSPrep(self.mps, new_wires)
+        new_wires = Wires(
+            [wire_map.get(wire, wire) for wire in self.hyperparameters["input_wires"]]
+        )
+        new_work_wires = Wires(
+            [wire_map.get(wire, wire) for wire in self.hyperparameters["work_wires"]]
+        )
+
+        return MPSPrep(
+            self.mps, new_wires, new_work_wires, self.hyperparameters["right_canonicalize"]
+        )
 
     @classmethod
     def _primitive_bind_call(cls, mps, wires, id=None):
@@ -188,6 +400,96 @@ class MPSPrep(Operation):
             return type.__call__(cls, mps=mps, wires=wires, id=id)  # pragma: no cover
         return cls._primitive.bind(*mps, wires=wires, id=id)
 
+    def decomposition(self):  # pylint: disable=arguments-differ
+        filtered_hyperparameters = {
+            key: value for key, value in self.hyperparameters.items() if key != "input_wires"
+        }
+        return self.compute_decomposition(
+            self.parameters, wires=self.hyperparameters["input_wires"], **filtered_hyperparameters
+        )
+
+    @staticmethod
+    def compute_decomposition(
+        mps, wires, work_wires, right_canonicalize=False
+    ):  # pylint: disable=arguments-differ
+        r"""Representation of the operator as a product of other operators.
+        The decomposition follows Eq. (23) in `arXiv:2310.18410 <https://arxiv.org/pdf/2310.18410>`_.
+
+        Args:
+            mps (list[Array]):  list of arrays of rank-3 and rank-2 tensors representing an MPS state as a
+                product of site matrices.
+
+            wires (Sequence[int]): wires that the template acts on. It should match the number of MPS tensors.
+            work_wires (Sequence[int]): list of extra qubits needed in the decomposition. If the maximum dimension
+                of the MPS tensors is `2^k``, then k ``work_wires`` will be needed. If no ``work_wires`` are given,
+                this operator can only be executed on the ``lightning.tensor`` device. Default is ``None``.
+
+            right_canonicalize (bool): Indicates whether a conversion to right-canonical form should be performed
+                to the mps. Default is ``False``.
+
+        Returns:
+            list[.Operator]: Decomposition of the operator
+        """
+
+        if work_wires is None:
+            raise ValueError("The qml.MPSPrep decomposition requires `work_wires` to be specified.")
+
+        max_bond_dimension = 0
+        for i in range(len(mps) - 1):
+            bond_dim = mps[i].shape[-1]
+            max_bond_dimension = max(max_bond_dimension, bond_dim)
+
+        if max_bond_dimension > 2 ** len(work_wires):
+            raise ValueError(
+                f"Incorrect number of `work_wires`. At least {int(qml.math.ceil(qml.math.log2(max_bond_dimension)))} `work_wires` must be provided."
+            )
+
+        ops = []
+        n_wires = len(work_wires) + 1
+
+        mps = mps.copy()
+
+        # Transform the MPS to ensure that the generated matrix is unitary
+        if right_canonicalize:
+            mps = right_canonicalize_mps(mps)
+
+        mps[0] = mps[0].reshape((1, *mps[0].shape))
+        mps[-1] = mps[-1].reshape((*mps[-1].shape, 1))
+
+        interface, dtype = qml.math.get_interface(mps[0]), mps[0].dtype
+
+        for i, Ai in enumerate(mps):
+
+            # Encode the tensor Ai in a unitary matrix following Eq.23 in https://arxiv.org/pdf/2310.18410
+            vectors = []
+            for column in Ai:
+                vector = qml.math.zeros(2**n_wires, like=interface, dtype=dtype)
+
+                if interface == "jax":
+                    vector = vector.at[: len(column[0])].set(column[0])
+                    vector = vector.at[
+                        2 ** (n_wires - 1) : 2 ** (n_wires - 1) + len(column[1])
+                    ].set(column[1])
+
+                else:
+                    vector[: len(column[0])] = column[0]
+                    vector[2 ** (n_wires - 1) : 2 ** (n_wires - 1) + len(column[1])] = column[1]
+
+                vectors.append(vector)
+
+            vectors = qml.math.stack(vectors).T
+            # The unitary is completed using QR decomposition
+            d, k = vectors.shape
+            new_columns = qml.math.array(np.random.RandomState(42).random((d, d - k)))
+            unitary_matrix, R = qml.math.linalg.qr(qml.math.hstack([vectors, new_columns]))
+            unitary_matrix *= qml.math.sign(
+                qml.math.diag(R)
+            )  # Enforce uniqueness for QR decomposition
+
+            ops.append(qml.QubitUnitary(unitary_matrix, wires=[wires[i]] + work_wires))
+
+        return ops
+
 
 if MPSPrep._primitive is not None:  # pylint: disable=protected-access
 


Test patch:
diff --git a/tests/templates/test_state_preparations/test_state_prep_mps.py b/tests/templates/test_state_preparations/test_state_prep_mps.py
index 7821c6ca0..20ba806a4 100644
--- a/tests/templates/test_state_preparations/test_state_prep_mps.py
+++ b/tests/templates/test_state_preparations/test_state_prep_mps.py
@@ -19,234 +19,636 @@ import numpy as np
 import pytest
 
 import pennylane as qml
+from pennylane.templates.state_preparations.state_prep_mps import (
+    _validate_mps_shape,
+    right_canonicalize_mps,
+)
 
 
-def test_standard_validity():
-    """Check the operation using the assert_valid function."""
-    mps = [
-        np.array([[0.0, 0.107], [0.994, 0.0]]),
-        np.array(
-            [
-                [[0.0, 0.0, 0.0, -0.0], [1.0, 0.0, 0.0, -0.0]],
-                [[0.0, 1.0, 0.0, -0.0], [0.0, 0.0, 0.0, -0.0]],
-            ]
-        ),
-        np.array(
-            [
-                [[-1.0, 0.0], [0.0, 0.0]],
-                [[0.0, 0.0], [0.0, 1.0]],
-                [[0.0, -1.0], [0.0, 0.0]],
-                [[0.0, 0.0], [1.0, 0.0]],
-            ]
-        ),
-        np.array([[-1.0, -0.0], [-0.0, -1.0]]),
-    ]
-    op = qml.MPSPrep(mps, wires=[0, 1, 2])
-    qml.ops.functions.assert_valid(op, skip_differentiation=True)
-
-
-def test_access_to_param():
-    """tests that the parameter is accessible."""
-    mps = [
-        np.array([[0.0, 0.107], [0.994, 0.0]]),
-        np.array(
-            [
-                [[0.0, 0.0, 0.0, -0.0], [1.0, 0.0, 0.0, -0.0]],
-                [[0.0, 1.0, 0.0, -0.0], [0.0, 0.0, 0.0, -0.0]],
-            ]
-        ),
-        np.array(
-            [
-                [[-1.0, 0.0], [0.0, 0.0]],
-                [[0.0, 0.0], [0.0, 1.0]],
-                [[0.0, -1.0], [0.0, 0.0]],
-                [[0.0, 0.0], [1.0, 0.0]],
-            ]
-        ),
-        np.array([[-1.0, -0.0], [-0.0, -1.0]]),
-    ]
-    op = qml.MPSPrep(mps, wires=[0, 1, 2])
-
-    for arr1, arr2 in zip(mps, op.mps):
-        assert np.allclose(arr1, arr2)
-
-
-@pytest.mark.parametrize(
-    ("mps", "msg_match"),
-    [
-        (
-            [
-                np.array([[0.0, 0.107, 0.0], [0.994, 0.0, 0.0]]),
-                np.array(
-                    [
-                        [[0.0, 0.0, 0.0, -0.0], [1.0, 0.0, 0.0, -0.0]],
-                        [[0.0, 1.0, 0.0, -0.0], [0.0, 0.0, 0.0, -0.0]],
-                    ]
-                ),
-                np.array(
-                    [
-                        [[-1.0, 0.0], [0.0, 0.0]],
-                        [[0.0, 0.0], [0.0, 1.0]],
-                        [[0.0, -1.0], [0.0, 0.0]],
-                        [[0.0, 0.0], [1.0, 0.0]],
-                    ]
-                ),
-                np.array([[-1.0, -0.0], [-0.0, -1.0]]),
-            ],
-            "The second dimension of the first tensor must be a power of 2.",
-        ),
-        (
-            [
-                np.array([0.0, 0.107, 0.994, 0.0]),
-                np.array(
-                    [
-                        [[0.0, 0.0, 0.0, -0.0], [1.0, 0.0, 0.0, -0.0]],
-                        [[0.0, 1.0, 0.0, -0.0], [0.0, 0.0, 0.0, -0.0]],
-                    ]
-                ),
-                np.array(
-                    [
-                        [[-1.0, 0.0], [0.0, 0.0]],
-                        [[0.0, 0.0], [0.0, 1.0]],
-                        [[0.0, -1.0], [0.0, 0.0]],
-                        [[0.0, 0.0], [1.0, 0.0]],
-                    ]
-                ),
-                np.array([[-1.0, -0.0], [-0.0, -1.0]]),
-            ],
-            "The first tensor must have exactly 2 dimensions",
-        ),
-        (
-            [
-                np.array([[0.0, 0.107], [0.994, 0.0]]),
-                np.array(
-                    [
-                        [[0.0, 0.0, 0.0, -0.0], [1.0, 0.0, 0.0, -0.0]],
-                        [[0.0, 1.0, 0.0, -0.0], [0.0, 0.0, 0.0, -0.0]],
-                    ]
-                ),
-                np.array(
+class TestMPSPrep:
+
+    def test_standard_validity(self):
+        """Check the template using the `assert_valid` function."""
+        mps = [
+            np.array([[0.0, 0.107], [0.994, 0.0]]),
+            np.array(
+                [
+                    [[0.0, 0.0, 0.0, -0.0], [1.0, 0.0, 0.0, -0.0]],
+                    [[0.0, 1.0, 0.0, -0.0], [0.0, 0.0, 0.0, -0.0]],
+                ]
+            ),
+            np.array(
+                [
                     [[-1.0, 0.0], [0.0, 0.0]],
-                ),
-                np.array([[-1.0, -0.0], [-0.0, -1.0]]),
-            ],
-            "Tensor 2 must have exactly 3 dimensions.",
-        ),
-        (
-            [
-                np.array([[0.0, 0.107], [0.994, 0.0]]),
-                np.array(
-                    [
-                        [[0.0, 0.0, 0.0, -0.0], [1.0, 0.0, 0.0, -0.0]],
-                        [[0.0, 1.0, 0.0, -0.0], [0.0, 0.0, 0.0, -0.0]],
-                    ]
-                ),
-                np.array(
-                    [
-                        [[-1.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
-                        [[0.0, 0.0, 0.0], [0.0, 1.0, 0.0]],
-                        [[0.0, -1.0, 0.0], [0.0, 0.0, 0.0]],
-                    ]
-                ),
-                np.array([[-1.0, -0.0], [-0.0, -1.0]]),
-            ],
-            "The first dimension of tensor 2 must be a power of 2.",
-        ),
-        (
-            [
-                np.array([[0.0, 0.107], [0.994, 0.0]]),
-                np.array(
-                    [
-                        [[-1.0, 0.0], [0.0, 0.0]],
-                        [[0.0, 0.0], [0.0, 1.0]],
-                        [[0.0, -1.0], [0.0, 0.0]],
-                        [[0.0, 0.0], [1.0, 0.0]],
-                    ]
-                ),
-                np.array(
-                    [
+                    [[0.0, 0.0], [0.0, 1.0]],
+                    [[0.0, -1.0], [0.0, 0.0]],
+                    [[0.0, 0.0], [1.0, 0.0]],
+                ]
+            ),
+            np.array([[-1.0, -0.0], [-0.0, -1.0]]),
+        ]
+
+        op = qml.MPSPrep(mps, wires=[0, 1, 2, 3], work_wires=[4, 5])
+        qml.ops.functions.assert_valid(op, skip_differentiation=True)
+
+    def test_access_to_param(self):
+        """Tests that the template parameters are accessible."""
+        mps = [
+            np.array([[0.0, 0.107], [0.994, 0.0]]),
+            np.array(
+                [
+                    [[0.0, 0.0, 0.0, -0.0], [1.0, 0.0, 0.0, -0.0]],
+                    [[0.0, 1.0, 0.0, -0.0], [0.0, 0.0, 0.0, -0.0]],
+                ]
+            ),
+            np.array(
+                [
+                    [[-1.0, 0.0], [0.0, 0.0]],
+                    [[0.0, 0.0], [0.0, 1.0]],
+                    [[0.0, -1.0], [0.0, 0.0]],
+                    [[0.0, 0.0], [1.0, 0.0]],
+                ]
+            ),
+            np.array([[-1.0, -0.0], [-0.0, -1.0]]),
+        ]
+        op = qml.MPSPrep(mps, wires=[0, 1, 2])
+
+        for arr1, arr2 in zip(mps, op.mps):
+            assert np.allclose(arr1, arr2)
+
+    @pytest.mark.parametrize(
+        ("mps", "msg_match"),
+        [
+            (
+                [
+                    np.array([[0.0, 0.107, 0.0], [0.994, 0.0, 0.0]]),
+                    np.array(
+                        [
+                            [[0.0, 0.0, 0.0, -0.0], [1.0, 0.0, 0.0, -0.0]],
+                            [[0.0, 1.0, 0.0, -0.0], [0.0, 0.0, 0.0, -0.0]],
+                        ]
+                    ),
+                    np.array(
+                        [
+                            [[-1.0, 0.0], [0.0, 0.0]],
+                            [[0.0, 0.0], [0.0, 1.0]],
+                            [[0.0, -1.0], [0.0, 0.0]],
+                            [[0.0, 0.0], [1.0, 0.0]],
+                        ]
+                    ),
+                    np.array([[-1.0, -0.0], [-0.0, -1.0]]),
+                ],
+                "The second dimension of the first tensor must be a power of 2.",
+            ),
+            (
+                [
+                    np.array([0.0, 0.107, 0.994, 0.0]),
+                    np.array(
+                        [
+                            [[0.0, 0.0, 0.0, -0.0], [1.0, 0.0, 0.0, -0.0]],
+                            [[0.0, 1.0, 0.0, -0.0], [0.0, 0.0, 0.0, -0.0]],
+                        ]
+                    ),
+                    np.array(
+                        [
+                            [[-1.0, 0.0], [0.0, 0.0]],
+                            [[0.0, 0.0], [0.0, 1.0]],
+                            [[0.0, -1.0], [0.0, 0.0]],
+                            [[0.0, 0.0], [1.0, 0.0]],
+                        ]
+                    ),
+                    np.array([[-1.0, -0.0], [-0.0, -1.0]]),
+                ],
+                "The first tensor must have exactly 2 dimensions",
+            ),
+            (
+                [
+                    np.array([[0.0, 0.107], [0.994, 0.0]]),
+                    np.array(
+                        [
+                            [[0.0, 0.0, 0.0, -0.0], [1.0, 0.0, 0.0, -0.0]],
+                            [[0.0, 1.0, 0.0, -0.0], [0.0, 0.0, 0.0, -0.0]],
+                        ]
+                    ),
+                    np.array(
                         [[-1.0, 0.0], [0.0, 0.0]],
-                        [[0.0, 0.0], [0.0, 1.0]],
-                        [[0.0, -1.0], [0.0, 0.0]],
-                        [[0.0, 0.0], [1.0, 0.0]],
-                    ]
-                ),
-                np.array([[-1.0, -0.0], [-0.0, -1.0]]),
-            ],
-            "Dimension mismatch:",
-        ),
-        (
-            [
-                np.array([[0.0, 0.107], [0.994, 0.0]]),
-                np.array(
-                    [
-                        [[0.0, 0.0, 0.0, -0.0], [1.0, 0.0, 0.0, -0.0]],
-                        [[0.0, 1.0, 0.0, -0.0], [0.0, 0.0, 0.0, -0.0]],
-                    ]
-                ),
+                    ),
+                    np.array([[-1.0, -0.0], [-0.0, -1.0]]),
+                ],
+                "Tensor 2 must have exactly 3 dimensions.",
+            ),
+            (
+                [
+                    np.array([[0.0, 0.107], [0.994, 0.0]]),
+                    np.array(
+                        [
+                            [[0.0, 0.0, 0.0, -0.0], [1.0, 0.0, 0.0, -0.0]],
+                            [[0.0, 1.0, 0.0, -0.0], [0.0, 0.0, 0.0, -0.0]],
+                        ]
+                    ),
+                    np.array(
+                        [
+                            [[-1.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
+                            [[0.0, 0.0, 0.0], [0.0, 1.0, 0.0]],
+                            [[0.0, -1.0, 0.0], [0.0, 0.0, 0.0]],
+                        ]
+                    ),
+                    np.array([[-1.0, -0.0], [-0.0, -1.0]]),
+                ],
+                "The first dimension of tensor 2 must be a power of 2.",
+            ),
+            (
+                [
+                    np.array([[0.0, 0.107], [0.994, 0.0]]),
+                    np.array(
+                        [
+                            [[-1.0, 0.0], [0.0, 0.0]],
+                            [[0.0, 0.0], [0.0, 1.0]],
+                            [[0.0, -1.0], [0.0, 0.0]],
+                            [[0.0, 0.0], [1.0, 0.0]],
+                        ]
+                    ),
+                    np.array(
+                        [
+                            [[-1.0, 0.0], [0.0, 0.0]],
+                            [[0.0, 0.0], [0.0, 1.0]],
+                            [[0.0, -1.0], [0.0, 0.0]],
+                            [[0.0, 0.0], [1.0, 0.0]],
+                        ]
+                    ),
+                    np.array([[-1.0, -0.0], [-0.0, -1.0]]),
+                ],
+                "Dimension mismatch:",
+            ),
+            (
+                [
+                    np.array([[0.0, 0.107], [0.994, 0.0]]),
+                    np.array(
+                        [
+                            [[0.0, 0.0, 0.0, -0.0], [1.0, 0.0, 0.0, -0.0]],
+                            [[0.0, 1.0, 0.0, -0.0], [0.0, 0.0, 0.0, -0.0]],
+                        ]
+                    ),
+                    np.array(
+                        [
+                            [[-1.0, 0.0], [0.0, 0.0]],
+                            [[0.0, 0.0], [0.0, 1.0]],
+                            [[0.0, -1.0], [0.0, 0.0]],
+                            [[0.0, 0.0], [1.0, 0.0]],
+                        ]
+                    ),
+                    np.array([-1.0, -0.0, -0.0, -1.0]),
+                ],
+                "The last tensor must have exactly 2 dimensions.",
+            ),
+            (
+                [
+                    np.array([[0.0, 0.107], [0.994, 0.0]]),
+                    np.array(
+                        [
+                            [[0.0, 0.0, 0.0, -0.0], [1.0, 0.0, 0.0, -0.0]],
+                            [[0.0, 1.0, 0.0, -0.0], [0.0, 0.0, 0.0, -0.0]],
+                        ]
+                    ),
+                    np.array(
+                        [
+                            [[-1.0, 0.0], [0.0, 0.0]],
+                            [[0.0, 0.0], [0.0, 1.0]],
+                            [[0.0, -1.0], [0.0, 0.0]],
+                            [[0.0, 0.0], [1.0, 0.0]],
+                        ]
+                    ),
+                    np.array([[-1.0, -0.0, 0.0], [-0.0, -1.0, 0.0]]),
+                ],
+                "The second dimension of the last tensor must be exactly 2.",
+            ),
+        ],
+    )
+    def test_mps_validate_function(self, mps, msg_match):
+        """Test that proper errors are raised in mps_validate_shape"""
+        with pytest.raises(AssertionError, match=msg_match):
+            _validate_mps_shape(mps)
+
+    @pytest.mark.jax
+    def test_jax_jit_mps(self):
+        """Check the operation works with jax and jit."""
+
+        import jax
+        from jax import numpy as jnp
+
+        mps = [
+            jnp.array([[0.0, 0.107j], [0.994, 0.0]]),
+            jnp.array(
+                [
+                    [[0.0, 0.0], [1.0, 0.0]],
+                    [[0.0, 1.0], [0.0, 0.0]],
+                ]
+            ),
+            jnp.array([[-1.0, -0.0], [-0.0, -1.0]]),
+        ]
+
+        state = [
+            0.0 + 0.0j,
+            -0.10705513j,
+            0.0 + 0.0j,
+            0.0 + 0.0j,
+            0.0 + 0.0j,
+            0.0 + 0.0j,
+            -0.99451217 + 0.0j,
+            0.0 + 0.0j,
+        ]
+
+        dev = qml.device("default.qubit")
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.MPSPrep(mps, wires=range(2, 5), work_wires=[0, 1])
+            return qml.state()
+
+        output = circuit()[:8]
+        output_jit = jax.jit(circuit)()[:8]
+
+        assert jax.numpy.allclose(output, jax.numpy.array(state), rtol=0.01)
+        assert jax.numpy.allclose(output_jit, jax.numpy.array(state), rtol=0.01)
+
+    @pytest.mark.parametrize(
+        ("mps", "state", "num_wires", "num_work_wires"),
+        [
+            (
+                [
+                    np.array([[0.70710678, 0.0], [0.0, 0.70710678]]),
+                    np.array(
+                        [
+                            [[0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0]],
+                            [[0.0, 0.0, -0.0, 0.0], [-1.0, 0.0, 0.0, 0.0]],
+                        ]
+                    ),
+                    np.array(
+                        [
+                            [[0.00000000e00, 1.74315280e-32], [-7.07106781e-01, -7.07106781e-01]],
+                            [[7.07106781e-01, 7.07106781e-01], [0.00000000e00, 0.00000000e00]],
+                            [[0.00000000e00, 0.00000000e00], [-7.07106781e-01, 7.07106781e-01]],
+                            [[-7.07106781e-01, 7.07106781e-01], [0.00000000e00, 0.00000000e00]],
+                        ]
+                    ),
+                    np.array([[1.0, 0.0], [0.0, 1.0]]),
+                ],
+                np.array([1 / 2, 1 / 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1 / 2, 1 / 2]),
+                4,
+                2,
+            ),
+            (
+                [
+                    np.array([[0.53849604, -0.44389787], [-0.59116842, -0.40434711]]),
+                    np.array(
+                        [
+                            [
+                                [-6.05052107e-01, 1.34284016e-01, -2.84018989e-01, -1.12416345e-01],
+                                [3.60988555e-01, 6.14571922e-01, -1.20681653e-01, -2.89527967e-04],
+                            ],
+                            [
+                                [-1.12393068e-01, 2.11496619e-01, 3.99193070e-01, -3.44891522e-01],
+                                [-5.99232567e-01, 3.76491467e-01, 3.04813277e-01, 2.65697349e-01],
+                            ],
+                        ]
+                    ),
+                    np.array(
+                        [
+                            [
+                                [0.87613189, -0.34254341, -0.12704983, -0.0161698],
+                                [-0.20758717, 0.1329479, -0.18184107, 0.06942658],
+                            ],
+                            [
+                                [-0.16499137, -0.13680142, -0.18432824, 0.12950892],
+                                [-0.68790868, -0.64141472, -0.12485688, -0.0556177],
+                            ],
+                            [
+                                [0.0352582, -0.37993402, 0.26781956, -0.25935129],
+                                [0.04351872, -0.27109361, 0.65111429, 0.4648453],
+                            ],
+                            [
+                                [0.1909576, 0.25461839, -0.07463641, -0.34390477],
+                                [-0.21279487, 0.0305474, 0.53420894, -0.66578494],
+                            ],
+                        ]
+                    ),
+                    np.array(
+                        [
+                            [[-0.26771292, -0.00628612], [-0.96316273, 0.02465422]],
+                            [[0.96011241, 0.07601506], [-0.2663889, 0.03798452]],
+                            [[-0.00727353, 0.4537835], [-0.02374101, -0.89076596]],
+                            [[0.08038064, -0.88784161], [-0.02812246, -0.45220057]],
+                        ]
+                    ),
+                    np.array([[-0.97855153, 0.2060022], [0.2060022, 0.97855153]]),
+                ],
                 np.array(
                     [
-                        [[-1.0, 0.0], [0.0, 0.0]],
-                        [[0.0, 0.0], [0.0, 1.0]],
-                        [[0.0, -1.0], [0.0, 0.0]],
-                        [[0.0, 0.0], [1.0, 0.0]],
+                        -0.30838839,
+                        0.00913474,
+                        -0.14842667,
+                        0.06058075,
+                        -0.00464704,
+                        0.16533888,
+                        0.08806463,
+                        0.19535211,
+                        0.17301488,
+                        -0.20745747,
+                        0.30696199,
+                        0.01131135,
+                        -0.08844845,
+                        -0.12525685,
+                        -0.13280471,
+                        0.29048878,
+                        0.19283544,
+                        0.02670539,
+                        0.39244675,
+                        0.00316343,
+                        -0.12347408,
+                        0.15527045,
+                        0.02744977,
+                        -0.01096892,
+                        -0.01303293,
+                        0.06124989,
+                        0.08509908,
+                        -0.08363092,
+                        -0.24638746,
+                        -0.00924703,
+                        0.44259161,
+                        -0.07738196,
                     ]
                 ),
-                np.array([-1.0, -0.0, -0.0, -1.0]),
-            ],
-            "The last tensor must have exactly 2 dimensions.",
-        ),
-        (
-            [
-                np.array([[0.0, 0.107], [0.994, 0.0]]),
+                5,
+                2,
+            ),
+            (
+                [
+                    np.array([[0.0, 0.107], [0.994, 0.0]]),
+                    np.array(
+                        [
+                            [[0.0, 0.0], [1.0, 0.0]],
+                            [[0.0, 1.0], [0.0, 0.0]],
+                        ]
+                    ),
+                    np.array([[-1.0, -0.0], [-0.0, -1.0]]),
+                ],
                 np.array(
                     [
-                        [[0.0, 0.0, 0.0, -0.0], [1.0, 0.0, 0.0, -0.0]],
-                        [[0.0, 1.0, 0.0, -0.0], [0.0, 0.0, 0.0, -0.0]],
+                        0.0 + 0.0j,
+                        -0.10705513 + 0.0j,
+                        0.0 + 0.0j,
+                        0.0 + 0.0j,
+                        0.0 + 0.0j,
+                        0.0 + 0.0j,
+                        -0.99451217 + 0.0j,
+                        0.0 + 0.0j,
                     ]
                 ),
+                3,
+                2,
+            ),
+            (
+                [
+                    np.array([[-0.6734052, 0.73927359], [0.73927359, 0.6734052]]),
+                    np.array(
+                        [
+                            [
+                                [-7.20105790e-01, -1.89741090e-01, 5.82838055e-01, -3.25185749e-01],
+                                [4.29632311e-01, -8.68380029e-01, 2.47651976e-01, -8.37507469e-04],
+                            ],
+                            [
+                                [
+                                    -1.00441887e-01,
+                                    -2.24394589e-01,
+                                    -6.15113611e-01,
+                                    -7.49128654e-01,
+                                ],
+                                [-5.35513918e-01, -3.99451514e-01, -4.69684517e-01, 5.77113322e-01],
+                            ],
+                        ]
+                    ),
+                    np.array(
+                        [
+                            [
+                                [-0.80632643, -0.47294461, 0.23182434, -0.03532423],
+                                [0.19104774, 0.18355921, 0.33180043, 0.15166789],
+                            ],
+                            [
+                                [-0.12789955, 0.15909324, -0.28329785, -0.23830545],
+                                [-0.5332595, 0.74593332, -0.19189508, 0.10234046],
+                            ],
+                            [
+                                [0.01881935, 0.30423266, 0.28341988, 0.32859376],
+                                [0.02322847, 0.21707857, 0.68904115, -0.58895129],
+                            ],
+                            [
+                                [-0.07230691, 0.14463917, 0.05603212, -0.3091061],
+                                [0.08057569, 0.01735283, -0.40104915, -0.59841618],
+                            ],
+                        ]
+                    ),
+                    np.array(
+                        [
+                            [[-0.22294664, -0.00953879], [-0.80210505, 0.03741124]],
+                            [[-0.53296709, -0.07688778], [0.1478749, -0.03842062]],
+                            [[-0.00305516, 0.34730904], [-0.00997211, -0.68175905]],
+                            [[-0.02820048, 0.56757052], [0.0098664, 0.28907832]],
+                        ]
+                    ),
+                    np.array([[-0.85785372, 0.18059319], [0.09911105, 0.47079729]]),
+                ],
                 np.array(
                     [
-                        [[-1.0, 0.0], [0.0, 0.0]],
-                        [[0.0, 0.0], [0.0, 1.0]],
-                        [[0.0, -1.0], [0.0, 0.0]],
-                        [[0.0, 0.0], [1.0, 0.0]],
+                        -0.30838839,
+                        0.00913474,
+                        -0.14842667,
+                        0.06058075,
+                        -0.00464704,
+                        0.16533888,
+                        0.08806463,
+                        0.19535211,
+                        0.17301488,
+                        -0.20745747,
+                        0.30696199,
+                        0.01131135,
+                        -0.08844845,
+                        -0.12525685,
+                        -0.13280471,
+                        0.29048878,
+                        0.19283544,
+                        0.02670539,
+                        0.39244675,
+                        0.00316343,
+                        -0.12347408,
+                        0.15527045,
+                        0.02744977,
+                        -0.01096892,
+                        -0.01303293,
+                        0.06124989,
+                        0.08509908,
+                        -0.08363092,
+                        -0.24638746,
+                        -0.00924703,
+                        0.44259161,
+                        -0.07738196,
                     ]
                 ),
-                np.array([[-1.0, -0.0, 0.0], [-0.0, -1.0, 0.0]]),
-            ],
-            "The second dimension of the last tensor must be exactly 2.",
-        ),
-    ],
-)
-def test_MPSPrep_error(mps, msg_match):
-    """Test that proper errors are raised for MPSPrep"""
-    with pytest.raises(AssertionError, match=msg_match):
-        qml.MPSPrep(mps, wires=[0, 1, 2])
-
-
-@pytest.mark.jax
-def test_jax_mps():
-    """Check the operation works with jax."""
-
-    from jax import numpy as jnp
-
-    mps = [
-        jnp.array([[0.0, 0.107], [0.994, 0.0]]),
-        jnp.array(
-            [
-                [[0.0, 0.0, 0.0, -0.0], [1.0, 0.0, 0.0, -0.0]],
-                [[0.0, 1.0, 0.0, -0.0], [0.0, 0.0, 0.0, -0.0]],
-            ]
-        ),
-        jnp.array(
-            [
-                [[-1.0, 0.0], [0.0, 0.0]],
-                [[0.0, 0.0], [0.0, 1.0]],
-                [[0.0, -1.0], [0.0, 0.0]],
-                [[0.0, 0.0], [1.0, 0.0]],
-            ]
-        ),
-        jnp.array([[-1.0, -0.0], [-0.0, -1.0]]),
-    ]
-    _ = qml.MPSPrep(mps, wires=[0, 1, 2])
+                5,
+                2,
+            ),
+        ],
+    )
+    def test_correctness(self, mps, state, num_wires, num_work_wires):
+        """Test correctness of the state approximation with MPS
+
+        Developer Note:
+        This data was generated by adapting the functionality `decompose_dense` in
+        the `pennylane-lightning/pennylane_lightning/lightning_tensor/_tensornet.py`
+
+        Args:
+            mps (List[Array]):  list of arrays of rank-3 and rank-2 tensors representing an MPS state as a
+                product of site matrices.
+            state (Array): target state that the mps is approaching
+            num_wires(int): number of wires to encode the state
+            num_work_wires(int): number of auxiliary wires used in the mps decomposition
+        """
+
+        wires = qml.registers({"work": num_work_wires, "state": num_wires})
+        dev = qml.device("default.qubit")
+
+        qs = qml.tape.QuantumScript(
+            qml.MPSPrep.compute_decomposition(
+                mps, wires=wires["state"], work_wires=wires["work"], right_canonicalize=True
+            ),
+            [qml.state()],
+        )
+        output = dev.execute(qs)[: 2**num_wires]
+
+        assert np.allclose(state, output, rtol=0.01)
+
+    def test_decomposition(self):
+        """Tests that the template defines the correct decomposition."""
+
+        mps = [
+            np.array([[0.70710678, 0.0], [0.0, 0.70710678]]),
+            np.array(
+                [
+                    [[0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0]],
+                    [[0.0, 0.0, -0.0, 0.0], [-1.0, 0.0, 0.0, 0.0]],
+                ]
+            ),
+            np.array(
+                [
+                    [[0.00000000e00, 1.74315280e-32], [-7.07106781e-01, -7.07106781e-01]],
+                    [[7.07106781e-01, 7.07106781e-01], [0.00000000e00, 0.00000000e00]],
+                    [[0.00000000e00, 0.00000000e00], [-7.07106781e-01, 7.07106781e-01]],
+                    [[-7.07106781e-01, 7.07106781e-01], [0.00000000e00, 0.00000000e00]],
+                ]
+            ),
+            np.array([[1.0, 0.0], [0.0, 1.0]]),
+        ]
+
+        ops = qml.MPSPrep.compute_decomposition(mps, wires=range(2, 6), work_wires=[0, 1])
+
+        for ind, op in enumerate(ops):
+            assert op.wires == qml.wires.Wires([2 + ind] + [0, 1])
+            assert op.name == "QubitUnitary"
+
+    @pytest.mark.parametrize(
+        ("work_wires", "msg"),
+        [(None, "The qml.MPSPrep decomposition requires"), (1, "Incorrect number of `work_wires`")],
+    )
+    def test_wires_decomposition(self, work_wires, msg):
+        """Checks that error is shown if no `work_wires` are given in decomposition"""
+
+        mps = [
+            np.array([[0.70710678, 0.0], [0.0, 0.70710678]]),
+            np.array(
+                [
+                    [[0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0]],
+                    [[0.0, 0.0, -0.0, 0.0], [-1.0, 0.0, 0.0, 0.0]],
+                ]
+            ),
+            np.array(
+                [
+                    [[0.00000000e00, 1.74315280e-32], [-7.07106781e-01, -7.07106781e-01]],
+                    [[7.07106781e-01, 7.07106781e-01], [0.00000000e00, 0.00000000e00]],
+                    [[0.00000000e00, 0.00000000e00], [-7.07106781e-01, 7.07106781e-01]],
+                    [[-7.07106781e-01, 7.07106781e-01], [0.00000000e00, 0.00000000e00]],
+                ]
+            ),
+            np.array([[1.0, 0.0], [0.0, 1.0]]),
+        ]
+
+        op = qml.MPSPrep(mps, wires=range(2, 5), work_wires=work_wires)
+        with pytest.raises(ValueError, match=msg):
+            op.decomposition()
+
+    def test_right_canonical(self):
+        """Checks that the function `right_canonicalize_mps` generates MPS in the correct format"""
+
+        n_sites = 4
+        mps = (
+            [np.ones((2, 4))]
+            + [np.ones((4, 2, 4)) for _ in range(1, n_sites - 1)]
+            + [np.ones((4, 2))]
+        )
+        mps_rc = right_canonicalize_mps(mps)
+
+        for i in range(1, n_sites - 1):
+            tensor = mps_rc[i]
+
+            # Right-canonical definition
+            contraction_matrix = np.tensordot(tensor, tensor.conj(), axes=([1, 2], [1, 2]))
+            assert np.allclose(contraction_matrix, np.eye(tensor.shape[0]))
+
+    @pytest.mark.jax
+    def test_right_canonical_jax_jit(self):
+        """Checks that the function `right_canonicalize_mps` works with JAX and JIT"""
+
+        import jax
+        from jax import numpy as jnp
+
+        n_sites = 4
+        mps = (
+            [jnp.ones((2, 4))]
+            + [jnp.ones((4, 2, 4)) for _ in range(1, n_sites - 1)]
+            + [jnp.ones((4, 2))]
+        )
+        mps_rc = jax.jit(right_canonicalize_mps)(mps)
+
+        for i in range(1, n_sites - 1):
+            tensor = mps_rc[i]
+
+            # Right-canonical definition
+            contraction_matrix = np.tensordot(tensor, tensor.conj(), axes=([1, 2], [1, 2]))
+            assert qml.math.allclose(contraction_matrix, qml.math.eye(tensor.shape[0]))
+
+    def test_immutable_input(self):
+        """Verifies that the input MPS remains unchanged after processing."""
+
+        n_sites = 4
+        mps = (
+            [np.ones((2, 4))]
+            + [np.ones((4, 2, 4)) for _ in range(1, n_sites - 1)]
+            + [np.ones((4, 2))]
+        )
+        mps_copy = mps.copy()
+        _ = right_canonicalize_mps(mps)
+
+        for tensor1, tensor2 in zip(mps, mps_copy):
+            assert qml.math.allclose(tensor1, tensor2)
+
+        wires = qml.registers({"work": 2, "state": n_sites})
+        dev = qml.device("default.qubit")
+
+        qs = qml.tape.QuantumScript(
+            qml.MPSPrep.compute_decomposition(
+                mps, wires=wires["state"], work_wires=wires["work"], right_canonicalize=True
+            ),
+            [qml.state()],
+        )
+        _ = dev.execute(qs)
+
+        for tensor1, tensor2 in zip(mps, mps_copy):
+            assert qml.math.allclose(tensor1, tensor2)