From dcfab7376d60136ea1f0496726dd23b9ea45d8ee Mon Sep 17 00:00:00 2001
From: Pietropaolo Frisoni <pietropaolo.frisoni@xanadu.ai>
Date: Mon, 29 Apr 2024 15:20:54 -0400
Subject: [PATCH 1/4] Creating the `LightningTensor` device class based on the
 new API (#671)

* empty commit (triggering CI)

* Auto update version

* Definition of the two front-end classes

* adding the `lightning_tensor` string to the supported backends in `setup.py`

* adding `__init__.py` file to directly import the `lightning_tensor` device class

* re-naming file

* Auto update version

* Creating the first prototype of initial MPS state tensor using `quimb`

* providing the `backend`, `method` parameters and making `wires` optional

* Changing names and structure

* Auto update version

* adding method required by the new device API design

* Auto update version

* using the `kwargs` parameter in `LightningTensor` and `CircuitMPS` in `quimb`

* taking some further inputs from the new device API

* Perhaps decided the overall structure of `LIghtningTensor`

* Auto update version

* adding docs to methods

* temporary changes so that `pylint` does not complain at this stage

* running `isort`

* re-running formatter after `isort`

* re-running formatter after `isort`

* Applying suggested formatting change from CI

* adding tmp unit tests

* Adding `quimb` in `requirements.txt`

* runing `isort` on mps test

* removing `quimb` from requirement and deleting unit tests for `lightning.tensor`

* Auto update version

* re-inserting unit tests with an additional `yml` file

* running isort on quimb test

* changing name of yml file

* preventing error in import

* updating yml file

* inserting `quimb` package in requirements-dev

* strange error with `quimb`

* strange error with `quimb`

* specifying scipy version

* removing installation of scipy from yml file

* removing the new `yml` file

* testing if tests are tested

* Covering all lines in tests

* forgot final line for formatter

* Python formatter on CI complaints

* covering missing lines

* formatter on CI complaints

* Trying not to skip test if Cpp is enabled

* skipping tests if Cpp is enabled

* removing the only line not covered by tests so far

* Auto update version

* Applying suggestions from code review and making the `state` attribute private (new API design)

* Python formatter

* removing params from `QuimbMPS`

* Auto update version

* removing `**kwargs` from `QuimbMPS`

* removing unnecessary param at this stage

* covering test line

* formatter...

* removing param description

* Making `pylint` happy

* forgot new arg in test

* Updating base class and `preprocess` function

* Updating `LightningTensor` class with new names from more advanced PR

* Auto update version

* Auto update version

* Triggering CI

* Auto update version

* Trying to remove pin from `quimb` in `requirements.dev`

* Auto update version

* Auto update version

* Removing infos on derivatives and using config options to pass parameters to interface

* Usual `pylint` failures

* Trying to solve formatting errors

* typo in docstring

* Sunday update: improved docstrings and structure

* Removing method that was supposed to be in next PR

* removing old TODO comment

* Removing changes from the `setup.py` file

* restoring previous format to `setup.py`

* Auto update version from '0.36.0-dev34' to '0.36.0-dev41'

* Auto update version from '0.36.0-dev40' to '0.36.0-dev41'

* Removing kwargs as suggested from code review

* Addressing comments from CR

* Skipping tests if CPP binary is available

* Auto update version from '0.36.0-dev42' to '0.36.0-dev43'

* Auto update version from '0.36.0-dev43' to '0.36.0-dev44'

* Restoring name in changelog (?)

* Increasing time limit for Python tests

* Applying suggestions from code review

---------

Co-authored-by: Dev version update bot <github-actions[bot]@users.noreply.github.com>
Co-authored-by: ringo-but-quantum <github-ringo-but-quantum@xanadu.ai>
---
 .github/CHANGELOG.md                          |   5 +-
 .github/workflows/tests_linux_python.yml      |   6 +-
 pennylane_lightning/core/_version.py          |   2 +-
 .../lightning_tensor/__init__.py              |  18 +
 .../lightning_tensor/backends/quimb/_mps.py   | 100 +++++
 .../lightning_tensor/lightning_tensor.py      | 348 ++++++++++++++++++
 requirements-dev.txt                          |   3 +-
 tests/lightning_tensor/__init__.py            |   0
 .../lightning_tensor/test_lightning_tensor.py | 115 ++++++
 tests/lightning_tensor/test_quimb_mps.py      |  51 +++
 10 files changed, 642 insertions(+), 6 deletions(-)
 create mode 100644 pennylane_lightning/lightning_tensor/__init__.py
 create mode 100644 pennylane_lightning/lightning_tensor/backends/quimb/_mps.py
 create mode 100644 pennylane_lightning/lightning_tensor/lightning_tensor.py
 create mode 100644 tests/lightning_tensor/__init__.py
 create mode 100644 tests/lightning_tensor/test_lightning_tensor.py
 create mode 100644 tests/lightning_tensor/test_quimb_mps.py

diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index b7aab6ea9e..fcba09ca10 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -2,6 +2,9 @@
 
 ### New features since last release
 
+* Add Python class for the `lightning.tensor` device which uses the new device API and the interface for `quimb` based on the MPS method.
+  [(#671)](https://github.com/PennyLaneAI/pennylane-lightning/pull/671)
+
 * Add compile-time support for AVX2/512 streaming operations in `lightning.qubit`.
   [(#664)](https://github.com/PennyLaneAI/pennylane-lightning/pull/664)
 
@@ -133,7 +136,7 @@
 
 This release contains contributions from (in alphabetical order):
 
-Ali Asadi, Amintor Dusko, Christina Lee, Vincent Michaud-Rioux, Lee James O'Riordan, Mudit Pandey, Shuli Shu
+Ali Asadi, Amintor Dusko, Pietropaolo Frisoni, Christina Lee, Vincent Michaud-Rioux, Lee James O'Riordan, Mudit Pandey, Shuli Shu
 
 ---
 
diff --git a/.github/workflows/tests_linux_python.yml b/.github/workflows/tests_linux_python.yml
index 5e6d9f0af9..4a61f7a70d 100644
--- a/.github/workflows/tests_linux_python.yml
+++ b/.github/workflows/tests_linux_python.yml
@@ -37,7 +37,7 @@ jobs:
       matrix:
         os: [ubuntu-22.04]
         pl_backend: ["lightning_qubit"]
-    timeout-minutes: 60
+    timeout-minutes: 75
     name: Python tests
     runs-on: ${{ matrix.os }}
 
@@ -144,7 +144,7 @@ jobs:
       matrix:
         os: [ubuntu-22.04]
         pl_backend: ["lightning_qubit"]
-    timeout-minutes: 60
+    timeout-minutes: 75
     name: Python tests with OpenBLAS
     runs-on: ${{ matrix.os }}
 
@@ -257,7 +257,7 @@ jobs:
         exclude:
           - pl_backend: ["all"]
             exec_model: OPENMP
-    timeout-minutes: 60
+    timeout-minutes: 75
     name: Python tests with Kokkos
     runs-on: ${{ matrix.os }}
 
diff --git a/pennylane_lightning/core/_version.py b/pennylane_lightning/core/_version.py
index 0dd7399263..b324311e95 100644
--- a/pennylane_lightning/core/_version.py
+++ b/pennylane_lightning/core/_version.py
@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.36.0-dev43"
+__version__ = "0.36.0-dev44"
diff --git a/pennylane_lightning/lightning_tensor/__init__.py b/pennylane_lightning/lightning_tensor/__init__.py
new file mode 100644
index 0000000000..48cc140c46
--- /dev/null
+++ b/pennylane_lightning/lightning_tensor/__init__.py
@@ -0,0 +1,18 @@
+# Copyright 2018-2024 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PennyLane lightning_tensor package."""
+
+from pennylane_lightning.core import __version__
+
+from .lightning_tensor import LightningTensor
diff --git a/pennylane_lightning/lightning_tensor/backends/quimb/_mps.py b/pennylane_lightning/lightning_tensor/backends/quimb/_mps.py
new file mode 100644
index 0000000000..c5f53b15ec
--- /dev/null
+++ b/pennylane_lightning/lightning_tensor/backends/quimb/_mps.py
@@ -0,0 +1,100 @@
+# Copyright 2018-2024 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Class implementation for the Quimb MPS interface for simulating quantum circuits while keeping the state always in MPS form.
+"""
+
+import numpy as np
+import pennylane as qml
+import quimb.tensor as qtn
+from pennylane.wires import Wires
+
+_operations = frozenset({})  # pragma: no cover
+# The set of supported operations.
+
+_observables = frozenset({})  # pragma: no cover
+# The set of supported observables.
+
+
+def stopping_condition(op: qml.operation.Operator) -> bool:
+    """A function that determines if an operation is supported by ``lightning.tensor`` for this interface."""
+    return op.name in _operations  # pragma: no cover
+
+
+def accepted_observables(obs: qml.operation.Operator) -> bool:
+    """A function that determines if an observable is supported by ``lightning.tensor`` for this interface."""
+    return obs.name in _observables  # pragma: no cover
+
+
+class QuimbMPS:
+    """Quimb MPS class.
+
+    Used internally by the `LightningTensor` device.
+    Interfaces with `quimb` for MPS manipulation, and provides methods to execute quantum circuits.
+
+    Args:
+        num_wires (int): the number of wires in the circuit.
+        interf_opts (dict): dictionary containing the interface options.
+        dtype (np.dtype): the complex type used for the MPS.
+    """
+
+    def __init__(self, num_wires, interf_opts, dtype=np.complex128):
+
+        if dtype not in [np.complex64, np.complex128]:  # pragma: no cover
+            raise TypeError(f"Unsupported complex type: {dtype}")
+
+        self._wires = Wires(range(num_wires))
+        self._dtype = dtype
+
+        self._init_state_ops = {
+            "binary": "0" * max(1, len(self._wires)),
+            "dtype": self._dtype.__name__,
+            "tags": [str(l) for l in self._wires.labels],
+        }
+
+        self._gate_opts = {
+            "contract": "swap+split",
+            "parametrize": None,
+            "cutoff": interf_opts["cutoff"],
+            "max_bond": interf_opts["max_bond_dim"],
+        }
+
+        self._expval_opts = {
+            "dtype": self._dtype.__name__,
+            "simplify_sequence": "ADCRS",
+            "simplify_atol": 0.0,
+        }
+
+        self._circuitMPS = qtn.CircuitMPS(psi0=self._initial_mps())
+
+    @property
+    def state(self):
+        """Current MPS handled by the interface."""
+        return self._circuitMPS.psi
+
+    def state_to_array(self) -> np.ndarray:
+        """Contract the MPS into a dense array."""
+        return self._circuitMPS.to_dense()
+
+    def _initial_mps(self) -> qtn.MatrixProductState:
+        r"""
+        Returns an initial state to :math:`\ket{0}`.
+
+        Internally, it uses `quimb`'s `MPS_computational_state` method.
+
+        Returns:
+            MatrixProductState: The initial MPS of a circuit.
+        """
+
+        return qtn.MPS_computational_state(**self._init_state_ops)
diff --git a/pennylane_lightning/lightning_tensor/lightning_tensor.py b/pennylane_lightning/lightning_tensor/lightning_tensor.py
new file mode 100644
index 0000000000..fc8974cda2
--- /dev/null
+++ b/pennylane_lightning/lightning_tensor/lightning_tensor.py
@@ -0,0 +1,348 @@
+# Copyright 2018-2024 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+This module contains the LightningTensor class that inherits from the new device interface.
+It is a device to perform tensor network simulation of a quantum circuit. 
+"""
+from dataclasses import replace
+from numbers import Number
+from typing import Callable, Optional, Sequence, Tuple, Union
+
+import numpy as np
+import pennylane as qml
+from pennylane.devices import DefaultExecutionConfig, Device, ExecutionConfig
+from pennylane.devices.modifiers import simulator_tracking, single_tape_support
+from pennylane.tape import QuantumTape
+from pennylane.transforms.core import TransformProgram
+from pennylane.typing import Result, ResultBatch
+
+from .backends.quimb._mps import QuimbMPS
+
+Result_or_ResultBatch = Union[Result, ResultBatch]
+QuantumTapeBatch = Sequence[QuantumTape]
+QuantumTape_or_Batch = Union[QuantumTape, QuantumTapeBatch]
+PostprocessingFn = Callable[[ResultBatch], Result_or_ResultBatch]
+
+
+_backends = frozenset({"quimb"})
+# The set of supported backends.
+
+_methods = frozenset({"mps"})
+# The set of supported methods.
+
+
+def accepted_backends(backend: str) -> bool:
+    """A function that determines whether or not a backend is supported by ``lightning.tensor``."""
+    return backend in _backends
+
+
+def accepted_methods(method: str) -> bool:
+    """A function that determines whether or not a method is supported by ``lightning.tensor``."""
+    return method in _methods
+
+
+@simulator_tracking
+@single_tape_support
+class LightningTensor(Device):
+    """PennyLane Lightning Tensor device.
+
+    A device to perform tensor network operations on a quantum circuit.
+
+    Args:
+        wires (int): The number of wires to initialize the device with.
+            Defaults to ``None`` if not specified.
+        backend (str): Supported backend. Currently, only ``quimb`` is supported.
+        method (str): Supported method. Currently, only ``mps`` is supported.
+        shots (int): How many times the circuit should be evaluated (or sampled) to estimate
+            the expectation values. Currently, it can only be ``None``, so that computation of
+            statistics like expectation values and variances is performed analytically.
+        c_dtype: Datatypes for the tensor representation. Must be one of
+            ``np.complex64`` or ``np.complex128``.
+        **kwargs: keyword arguments. The following options are currently supported:
+
+            ``max_bond_dim`` (int): Maximum bond dimension for the MPS simulator.
+                It corresponds to the number of Schmidt coefficients retained at the end of the SVD algorithm when applying gates. Default is ``None``.
+            ``cutoff`` (float): Truncation threshold for the Schmidt coefficients in a MPS simulator. Default is ``1e-16``.
+    """
+
+    # pylint: disable=too-many-instance-attributes
+
+    # So far we just consider the options for MPS simulator
+    _device_options = (
+        "backend",
+        "c_dtype",
+        "cutoff",
+        "method",
+        "max_bond_dim",
+    )
+
+    _new_API = True
+
+    # pylint: disable=too-many-arguments
+    def __init__(
+        self,
+        *,
+        wires=None,
+        backend="quimb",
+        method="mps",
+        shots=None,
+        c_dtype=np.complex128,
+        **kwargs,
+    ):
+
+        if not accepted_backends(backend):
+            raise ValueError(f"Unsupported backend: {backend}")
+
+        if not accepted_methods(method):
+            raise ValueError(f"Unsupported method: {method}")
+
+        if shots is not None:
+            raise ValueError("LightningTensor does not support finite shots.")
+
+        super().__init__(wires=wires, shots=shots)
+
+        self._num_wires = len(self.wires) if self.wires else 0
+        self._backend = backend
+        self._method = method
+        self._c_dtype = c_dtype
+
+        # options for MPS
+        self._max_bond_dim = kwargs.get("max_bond_dim", None)
+        self._cutoff = kwargs.get("cutoff", np.finfo(self._c_dtype).eps)
+
+        self._interface = None
+        interface_opts = self._setup_execution_config().device_options
+
+        if self.backend == "quimb" and self.method == "mps":
+            self._interface = QuimbMPS(
+                self._num_wires,
+                interface_opts,
+                self._c_dtype,
+            )
+
+        else:
+            raise ValueError(
+                f"Unsupported backend: {self.backend} or method: {self.method}"
+            )  # pragma: no cover
+
+        for arg in kwargs:
+            if arg not in self._device_options:
+                raise TypeError(
+                    f"Unexpected argument: {arg} during initialization of the LightningTensor device."
+                )
+
+    @property
+    def name(self):
+        """The name of the device."""
+        return "lightning.tensor"
+
+    @property
+    def num_wires(self):
+        """Number of wires addressed on this device."""
+        return self._num_wires
+
+    @property
+    def backend(self):
+        """Supported backend."""
+        return self._backend
+
+    @property
+    def method(self):
+        """Supported method."""
+        return self._method
+
+    @property
+    def c_dtype(self):
+        """Tensor complex data type."""
+        return self._c_dtype
+
+    dtype = c_dtype
+
+    def _setup_execution_config(
+        self, config: Optional[ExecutionConfig] = DefaultExecutionConfig
+    ) -> ExecutionConfig:
+        """
+        Update the execution config with choices for how the device should be used and the device options.
+        """
+        # TODO: add options for gradients next quarter
+
+        updated_values = {}
+
+        new_device_options = dict(config.device_options)
+        for option in self._device_options:
+            if option not in new_device_options:
+                new_device_options[option] = getattr(self, f"_{option}", None)
+
+        return replace(config, **updated_values, device_options=new_device_options)
+
+    def preprocess(
+        self,
+        execution_config: ExecutionConfig = DefaultExecutionConfig,
+    ):
+        """This function defines the device transform program to be applied and an updated device configuration.
+
+        Args:
+            execution_config (Union[ExecutionConfig, Sequence[ExecutionConfig]]): A data structure describing the
+                parameters needed to fully describe the execution.
+
+        Returns:
+            TransformProgram, ExecutionConfig: A transform program that when called returns :class:`~.QuantumTape`'s that the
+            device can natively execute as well as a postprocessing function to be called after execution, and a configuration
+            with unset specifications filled in.
+
+        This device:
+
+        * Supports any qubit operations that provide a matrix.
+        * Currently does not support finite shots.
+        """
+
+        config = self._setup_execution_config(execution_config)
+
+        program = TransformProgram()
+
+        # more in the next PR
+
+        return program, config
+
+    def execute(
+        self,
+        circuits: QuantumTape_or_Batch,
+        execution_config: ExecutionConfig = DefaultExecutionConfig,
+    ) -> Result_or_ResultBatch:
+        """Execute a circuit or a batch of circuits and turn it into results.
+
+        Args:
+            circuits (Union[QuantumTape, Sequence[QuantumTape]]): the quantum circuits to be executed.
+            execution_config (ExecutionConfig): a datastructure with additional information required for execution.
+
+        Returns:
+            TensorLike, tuple[TensorLike], tuple[tuple[TensorLike]]: A numeric result of the computation.
+        """
+        # comment is removed in the next PR
+        # return self._interface.execute(circuits, execution_config)
+
+    # pylint: disable=unused-argument
+    def supports_derivatives(
+        self,
+        execution_config: Optional[ExecutionConfig] = None,
+        circuit: Optional[qml.tape.QuantumTape] = None,
+    ) -> bool:
+        """Check whether or not derivatives are available for a given configuration and circuit.
+
+        Args:
+            execution_config (ExecutionConfig): The configuration of the desired derivative calculation.
+            circuit (QuantumTape): An optional circuit to check derivatives support for.
+
+        Returns:
+            Bool: Whether or not a derivative can be calculated provided the given information.
+
+        """
+        return False
+
+    def compute_derivatives(
+        self,
+        circuits: QuantumTape_or_Batch,
+        execution_config: ExecutionConfig = DefaultExecutionConfig,
+    ):
+        """Calculate the jacobian of either a single or a batch of circuits on the device.
+
+        Args:
+            circuits (Union[QuantumTape, Sequence[QuantumTape]]): the circuits to calculate derivatives for.
+            execution_config (ExecutionConfig): a datastructure with all additional information required for execution.
+
+        Returns:
+            Tuple: The jacobian for each trainable parameter.
+        """
+        raise NotImplementedError(
+            "The computation of derivatives has yet to be implemented for the lightning.tensor device."
+        )
+
+    def execute_and_compute_derivatives(
+        self,
+        circuits: QuantumTape_or_Batch,
+        execution_config: ExecutionConfig = DefaultExecutionConfig,
+    ):
+        """Compute the results and jacobians of circuits at the same time.
+
+        Args:
+            circuits (Union[QuantumTape, Sequence[QuantumTape]]): the circuits or batch of circuits.
+            execution_config (ExecutionConfig): a datastructure with all additional information required for execution.
+
+        Returns:
+            tuple: A numeric result of the computation and the gradient.
+        """
+        raise NotImplementedError(
+            "The computation of derivatives has yet to be implemented for the lightning.tensor device."
+        )
+
+    # pylint: disable=unused-argument
+    def supports_vjp(
+        self,
+        execution_config: Optional[ExecutionConfig] = None,
+        circuit: Optional[QuantumTape] = None,
+    ) -> bool:
+        """Whether or not this device defines a custom vector jacobian product.
+
+        Args:
+            execution_config (ExecutionConfig): The configuration of the desired derivative calculation.
+            circuit (QuantumTape): An optional circuit to check derivatives support for.
+
+        Returns:
+            Bool: Whether or not a derivative can be calculated provided the given information.
+        """
+        # TODO: implement during next quarter
+        return False
+
+    def compute_vjp(
+        self,
+        circuits: QuantumTape_or_Batch,
+        cotangents: Tuple[Number],
+        execution_config: ExecutionConfig = DefaultExecutionConfig,
+    ):
+        r"""The vector jacobian product used in reverse-mode differentiation.
+
+        Args:
+            circuits (Union[QuantumTape, Sequence[QuantumTape]]): the circuit or batch of circuits.
+            cotangents (Tuple[Number, Tuple[Number]]): Gradient-output vector. Must have shape matching the output shape of the
+                corresponding circuit. If the circuit has a single output, ``cotangents`` may be a single number, not an iterable
+                of numbers.
+            execution_config (ExecutionConfig): a datastructure with all additional information required for execution.
+
+        Returns:
+            tensor-like: A numeric result of computing the vector jacobian product.
+        """
+        raise NotImplementedError(
+            "The computation of vector jacobian product has yet to be implemented for the lightning.tensor device."
+        )
+
+    def execute_and_compute_vjp(
+        self,
+        circuits: QuantumTape_or_Batch,
+        cotangents: Tuple[Number],
+        execution_config: ExecutionConfig = DefaultExecutionConfig,
+    ):
+        """Calculate both the results and the vector jacobian product used in reverse-mode differentiation.
+
+        Args:
+            circuits (Union[QuantumTape, Sequence[QuantumTape]]): the circuit or batch of circuits to be executed.
+            cotangents (Tuple[Number, Tuple[Number]]): Gradient-output vector. Must have shape matching the output shape of the
+                corresponding circuit.
+            execution_config (ExecutionConfig): a datastructure with all additional information required for execution.
+
+        Returns:
+            Tuple, Tuple: the result of executing the scripts and the numeric result of computing the vector jacobian product
+        """
+        raise NotImplementedError(
+            "The computation of vector jacobian product has yet to be implemented for the lightning.tensor device."
+        )
diff --git a/requirements-dev.txt b/requirements-dev.txt
index c87c9154a4..9d0e09b02c 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -17,4 +17,5 @@ click==8.0.4
 cmake
 custatevec-cu12
 pylint
-scipy
+scipy~=1.12.0
+quimb
diff --git a/tests/lightning_tensor/__init__.py b/tests/lightning_tensor/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/lightning_tensor/test_lightning_tensor.py b/tests/lightning_tensor/test_lightning_tensor.py
new file mode 100644
index 0000000000..3d7d1033b8
--- /dev/null
+++ b/tests/lightning_tensor/test_lightning_tensor.py
@@ -0,0 +1,115 @@
+# Copyright 2018-2024 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Unit tests for the generic lightning tensor class.
+"""
+
+
+import numpy as np
+import pennylane as qml
+import pytest
+from conftest import LightningDevice  # tested device
+from pennylane.wires import Wires
+
+from pennylane_lightning.lightning_tensor import LightningTensor
+
+if not LightningDevice._new_API:
+    pytest.skip("Exclusive tests for new API. Skipping.", allow_module_level=True)
+
+if LightningDevice._CPP_BINARY_AVAILABLE:
+    pytest.skip("Device doesn't have C++ support yet.", allow_module_level=True)
+
+
+@pytest.mark.parametrize("num_wires", [None, 4])
+@pytest.mark.parametrize("c_dtype", [np.complex64, np.complex128])
+def test_device_name_and_init(num_wires, c_dtype):
+    """Test the class initialization and returned properties."""
+    wires = Wires(range(num_wires)) if num_wires else None
+    dev = LightningTensor(wires=wires, c_dtype=c_dtype)
+    assert dev.name == "lightning.tensor"
+    assert dev.c_dtype == c_dtype
+    assert dev.wires == wires
+    if num_wires is None:
+        assert dev.num_wires == 0
+    else:
+        assert dev.num_wires == num_wires
+
+
+@pytest.mark.parametrize("backend", ["fake_backend"])
+def test_invalid_backend(backend):
+    """Test an invalid backend."""
+    with pytest.raises(ValueError, match=f"Unsupported backend: {backend}"):
+        LightningTensor(backend=backend)
+
+
+@pytest.mark.parametrize("method", ["fake_method"])
+def test_invalid_method(method):
+    """Test an invalid method."""
+    with pytest.raises(ValueError, match=f"Unsupported method: {method}"):
+        LightningTensor(method=method)
+
+
+def test_invalid_keyword_arg():
+    """Test an invalid keyword argument."""
+    with pytest.raises(
+        TypeError,
+        match=f"Unexpected argument: fake_arg during initialization of the LightningTensor device.",
+    ):
+        LightningTensor(fake_arg=None)
+
+
+def test_invalid_shots():
+    """Test that an error is raised if finite number of shots are requestd."""
+    with pytest.raises(ValueError, match="LightningTensor does not support finite shots."):
+        LightningTensor(shots=5)
+
+
+def test_support_derivatives():
+    """Test that the device does not support derivatives yet."""
+    dev = LightningTensor()
+    assert not dev.supports_derivatives()
+
+
+def test_compute_derivatives():
+    """Test that an error is raised if the `compute_derivatives` method is called."""
+    dev = LightningTensor()
+    with pytest.raises(NotImplementedError):
+        dev.compute_derivatives(circuits=None)
+
+
+def test_execute_and_compute_derivatives():
+    """Test that an error is raised if `execute_and_compute_derivative` method is called."""
+    dev = LightningTensor()
+    with pytest.raises(NotImplementedError):
+        dev.execute_and_compute_derivatives(circuits=None)
+
+
+def test_supports_vjp():
+    """Test that the device does not support VJP yet."""
+    dev = LightningTensor()
+    assert not dev.supports_vjp()
+
+
+def test_compute_vjp():
+    """Test that an error is raised if `compute_vjp` method is called."""
+    dev = LightningTensor()
+    with pytest.raises(NotImplementedError):
+        dev.compute_vjp(circuits=None, cotangents=None)
+
+
+def test_execute_and_compute_vjp():
+    """Test that an error is raised if `execute_and_compute_vjp` method is called."""
+    dev = LightningTensor()
+    with pytest.raises(NotImplementedError):
+        dev.execute_and_compute_vjp(circuits=None, cotangents=None)
diff --git a/tests/lightning_tensor/test_quimb_mps.py b/tests/lightning_tensor/test_quimb_mps.py
new file mode 100644
index 0000000000..8f3e376e56
--- /dev/null
+++ b/tests/lightning_tensor/test_quimb_mps.py
@@ -0,0 +1,51 @@
+# Copyright 2018-2024 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Unit tests for the ``quimb`` interface.
+"""
+
+
+import numpy as np
+import pytest
+import quimb.tensor as qtn
+from conftest import LightningDevice  # tested device
+from pennylane.wires import Wires
+
+from pennylane_lightning.lightning_tensor import LightningTensor
+
+if not LightningDevice._new_API:
+    pytest.skip("Exclusive tests for new API. Skipping.", allow_module_level=True)
+
+if LightningDevice._CPP_BINARY_AVAILABLE:
+    pytest.skip("Device doesn't have C++ support yet.", allow_module_level=True)
+
+
+@pytest.mark.parametrize("backend", ["quimb"])
+@pytest.mark.parametrize("method", ["mps"])
+class TestQuimbMPS:
+    """Tests for the MPS method."""
+
+    @pytest.mark.parametrize("num_wires", [None, 4])
+    @pytest.mark.parametrize("c_dtype", [np.complex64, np.complex128])
+    def test_device_init(self, num_wires, c_dtype, backend, method):
+        """Test the class initialization and returned properties."""
+
+        wires = Wires(range(num_wires)) if num_wires else None
+        dev = LightningTensor(wires=wires, backend=backend, method=method, c_dtype=c_dtype)
+        assert isinstance(dev._interface.state, qtn.MatrixProductState)
+        assert isinstance(dev._interface.state_to_array(), np.ndarray)
+
+        _, config = dev.preprocess()
+        assert config.device_options["backend"] == backend
+        assert config.device_options["method"] == method

From 05a98aead121db6efdbd2c5cceff36ce8e1a9a4f Mon Sep 17 00:00:00 2001
From: Amintor Dusko <87949283+AmintorDusko@users.noreply.github.com>
Date: Mon, 29 Apr 2024 16:34:28 -0400
Subject: [PATCH 2/4] Increase tolerance for a flaky test
 (test_single_return_value) (#703)

* update dev version

* increase tolerance

* update changelog

* trigger CIs

* Auto update version from '0.36.0-dev44' to '0.36.0-dev45'

---------

Co-authored-by: ringo-but-quantum <github-ringo-but-quantum@xanadu.ai>
---
 .github/CHANGELOG.md                             | 3 +++
 pennylane_lightning/core/_version.py             | 2 +-
 tests/lightning_qubit/test_measurements_class.py | 2 +-
 3 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index fcba09ca10..700adc8e5b 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -132,6 +132,9 @@
 * Update the version of `codecov-action` to v4 and fix the CodeCov issue with the PL-Lightning check-compatibility actions.
   [(#682)](https://github.com/PennyLaneAI/pennylane-lightning/pull/682)
 
+* Increase tolerance for a flaky test.
+  [(#703)](https://github.com/PennyLaneAI/pennylane-lightning/pull/703)
+
 ### Contributors
 
 This release contains contributions from (in alphabetical order):
diff --git a/pennylane_lightning/core/_version.py b/pennylane_lightning/core/_version.py
index b324311e95..0198b9a3ca 100644
--- a/pennylane_lightning/core/_version.py
+++ b/pennylane_lightning/core/_version.py
@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.36.0-dev44"
+__version__ = "0.36.0-dev45"
diff --git a/tests/lightning_qubit/test_measurements_class.py b/tests/lightning_qubit/test_measurements_class.py
index bd070629f8..c282e69358 100644
--- a/tests/lightning_qubit/test_measurements_class.py
+++ b/tests/lightning_qubit/test_measurements_class.py
@@ -474,7 +474,7 @@ def test_single_return_value(self, measurement, observable, lightning_sv, tol):
         result = m.measure_final_state(tape)
 
         # a few tests may fail in single precision, and hence we increase the tolerance
-        assert np.allclose(result, expected, max(tol, 1.0e-5))
+        assert np.allclose(result, expected, max(tol, 1.0e-4))
 
     @flaky(max_runs=5)
     @pytest.mark.parametrize("shots", [None, 1000000])

From fdb47f072c0f1eff01ecbebfb17e21fb9b39f9a9 Mon Sep 17 00:00:00 2001
From: Rashid N H M <95639609+rashidnhm@users.noreply.github.com>
Date: Fri, 3 May 2024 09:28:02 -0400
Subject: [PATCH 3/4] On-Board Lightning Linux workflows (that are run on PRs)
 to large runners based on label (#700)

* Try on-boarding to large runner

* Auto update version from '0.36.0-dev43' to '0.36.0-dev44'

* Remove usage of runner.os in job name

* Fix invalid runs-on

* Remove usage of runner.os in job name

* Trigger CI

* Fix inputs.run-on for build&cache kokkos

* Trigger CI

* Auto update version from '0.36.0-dev45' to '0.36.0-dev46'

---------

Co-authored-by: ringo-but-quantum <github-ringo-but-quantum@xanadu.ai>
Co-authored-by: Vincent Michaud-Rioux <vincentm@nanoacademic.com>
---
 .../build_and_cache_Kokkos_linux.yml          | 10 ++++-
 .../workflows/determine-workflow-runner.yml   | 40 +++++++++++++++++++
 .github/workflows/tests_lgpu_cpp.yml          |  5 +--
 .github/workflows/tests_lgpu_python.yml       |  5 +--
 .github/workflows/tests_linux_cpp.yml         | 24 +++++++----
 .github/workflows/tests_linux_python.yml      | 20 +++++++---
 .github/workflows/tests_without_binary.yml    | 10 ++++-
 .github/workflows/wheel_linux_x86_64.yml      | 18 +++++----
 .github/workflows/wheel_linux_x86_64_cuda.yml | 14 +++++--
 .github/workflows/wheel_noarch.yml            | 13 ++++--
 pennylane_lightning/core/_version.py          |  2 +-
 11 files changed, 121 insertions(+), 40 deletions(-)
 create mode 100644 .github/workflows/determine-workflow-runner.yml

diff --git a/.github/workflows/build_and_cache_Kokkos_linux.yml b/.github/workflows/build_and_cache_Kokkos_linux.yml
index 75754c5a7e..e612d17a6b 100644
--- a/.github/workflows/build_and_cache_Kokkos_linux.yml
+++ b/.github/workflows/build_and_cache_Kokkos_linux.yml
@@ -9,6 +9,14 @@ on:
       os:
         required: true
         type: string
+      runs_on:
+        description: |
+          The runner that should run the jobs. If left blank, the value from inputs.os is used.
+          This is useful if you want the jobs to run in a specific runner group, while not using that group name as part 
+          of the cache key.
+        required: false
+        type: string
+        default: ''
     outputs:
       exec_model:
         description: "The execution model for Kokkos."
@@ -43,7 +51,7 @@ jobs:
         kokkos_version: ${{ fromJson(needs.linux-set-builder-matrix.outputs.kokkos_version) }}
     timeout-minutes: 30
     name: Kokkos core (${{ matrix.exec_model }})
-    runs-on: ${{ inputs.os }}
+    runs-on: ${{ inputs.runs_on || inputs.os }}
 
     steps:
       - name: Cache installation directories
diff --git a/.github/workflows/determine-workflow-runner.yml b/.github/workflows/determine-workflow-runner.yml
new file mode 100644
index 0000000000..78919a9762
--- /dev/null
+++ b/.github/workflows/determine-workflow-runner.yml
@@ -0,0 +1,40 @@
+name: Determine Workflow Runner group
+
+on:
+  workflow_call:
+    inputs:
+      default_runner:
+        description: The runner type that is used by the calling workflow by default
+        required: true
+        type: string
+    outputs:
+      runner_group:
+        description: The runner all subsequent jobs within the calling workflow should run on
+        value: ${{ jobs.determine_workflow_runner.outputs.runner_group || inputs.default_runner }}
+
+env:
+  LARGE_RUNNER_GROUP_NAME: pl-4-core-large-runner
+
+jobs:
+  determine_workflow_runner:
+    runs-on: >-
+      ${{
+        (
+          github.event_name == 'pull_request'
+          && contains(github.event.pull_request.labels.*.name, 'urgent')
+        ) && 'pl-4-core-large-runner' || 'ubuntu-latest'
+      }}
+
+    outputs:
+      runner_group: ${{ steps.runner_group.outputs.runner_group }}
+
+    steps:
+      - name: Output Runner Group name
+        if: >-
+          ${{
+            github.event_name == 'pull_request'
+            && contains(github.event.pull_request.labels.*.name, 'urgent')
+            && startsWith(inputs.default_runner, 'ubuntu')
+          }}
+        id: runner_group
+        run: echo "runner_group=$LARGE_RUNNER_GROUP_NAME" >> $GITHUB_OUTPUT
\ No newline at end of file
diff --git a/.github/workflows/tests_lgpu_cpp.yml b/.github/workflows/tests_lgpu_cpp.yml
index d50a56ce3e..9cdc46c389 100644
--- a/.github/workflows/tests_lgpu_cpp.yml
+++ b/.github/workflows/tests_lgpu_cpp.yml
@@ -42,8 +42,6 @@ jobs:
     strategy:
       max-parallel: 1
       matrix:
-        os: [ubuntu-22.04]
-        pl_backend: ["lightning_gpu"]
         cuda_version: ["12"]
 
     steps:
@@ -62,13 +60,12 @@ jobs:
     needs: [builddeps]
     strategy:
       matrix:
-        os: [ubuntu-22.04]
         pl_backend: ["lightning_gpu"]
         cuda_version: ["12"]
 
     name: C++ tests (Lightning-GPU)
     runs-on:
-      - ${{ matrix.os }}
+      - ubuntu-22.04
       - self-hosted
       - gpu
 
diff --git a/.github/workflows/tests_lgpu_python.yml b/.github/workflows/tests_lgpu_python.yml
index a7eb06260d..2e6c6482cb 100644
--- a/.github/workflows/tests_lgpu_python.yml
+++ b/.github/workflows/tests_lgpu_python.yml
@@ -42,8 +42,6 @@ jobs:
     strategy:
       max-parallel: 1
       matrix:
-        os: [ubuntu-22.04]
-        pl_backend: ["lightning_gpu"]
         cuda_version: ["12"]
 
     steps:
@@ -61,14 +59,13 @@ jobs:
     needs: [builddeps]
     strategy:
       matrix:
-        os: [ubuntu-22.04]
         pl_backend: ["lightning_gpu"]
         default_backend: ["lightning_qubit"]
         cuda_version: ["12"]
 
     name: Python tests with LGPU
     runs-on:
-      - ${{ matrix.os }}
+      - ubuntu-22.04
       - self-hosted
       - gpu
 
diff --git a/.github/workflows/tests_linux_cpp.yml b/.github/workflows/tests_linux_cpp.yml
index 2a4d2c7058..f0c33f80da 100644
--- a/.github/workflows/tests_linux_cpp.yml
+++ b/.github/workflows/tests_linux_cpp.yml
@@ -33,15 +33,21 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
+  determine_runner:
+    name: Determine runner type to use
+    uses: ./.github/workflows/determine-workflow-runner.yml
+    with:
+      default_runner: ubuntu-22.04
+
   cpptests:
     if: ${{ !contains(fromJSON('["schedule", "workflow_dispatch"]'), github.event_name) }}
+    needs: [determine_runner]
     strategy:
       matrix:
-        os: [ubuntu-22.04]
         pl_backend: ["lightning_qubit"]
     timeout-minutes: 60
     name: C++ tests
-    runs-on: ${{ matrix.os }}
+    runs-on: ${{ needs.determine_runner.outputs.runner_group }}
 
     steps:
       - uses: actions/setup-python@v5
@@ -123,13 +129,13 @@ jobs:
 
   cpptestswithOpenBLAS:
     if: ${{ !contains(fromJSON('["schedule", "workflow_dispatch"]'), github.event_name) }}
+    needs: [determine_runner]
     strategy:
       matrix:
-        os: [ubuntu-22.04]
         pl_backend: ["lightning_qubit"]
     timeout-minutes: 60
     name: C++ tests (OpenBLAS)
-    runs-on: ${{ matrix.os }}
+    runs-on: ${{ needs.determine_runner.outputs.runner_group }}
 
     steps:
       - uses: actions/setup-python@v5
@@ -180,13 +186,15 @@ jobs:
 
   build_and_cache_Kokkos:
     name: "Build and cache Kokkos"
+    needs: [determine_runner]
     uses: ./.github/workflows/build_and_cache_Kokkos_linux.yml
     with:
+      runs_on: ${{ needs.determine_runner.outputs.runner_group }}
       os: ubuntu-22.04
 
   cpptestswithKokkos:
     if: ${{ !contains(fromJSON('["schedule", "workflow_dispatch"]'), github.event_name) }}
-    needs: [build_and_cache_Kokkos]
+    needs: [build_and_cache_Kokkos, determine_runner]
     strategy:
       matrix:
         os: [ubuntu-22.04]
@@ -195,7 +203,7 @@ jobs:
         kokkos_version: ${{ fromJson(needs.build_and_cache_Kokkos.outputs.kokkos_version) }}
     timeout-minutes: 60
     name: C++ tests (Kokkos)
-    runs-on: ${{ matrix.os }}
+    runs-on: ${{ needs.determine_runner.outputs.runner_group }}
 
     steps:
       - uses: actions/setup-python@v5
@@ -281,7 +289,7 @@ jobs:
   cpptestsWithMultipleBackends:
   # Device-specific tests are performed for both. Device-agnostic tests default to LightningQubit.
     if: ${{ !contains(fromJSON('["schedule", "workflow_dispatch"]'), github.event_name) }}
-    needs: [build_and_cache_Kokkos]
+    needs: [build_and_cache_Kokkos, determine_runner]
     strategy:
       matrix:
         os: [ubuntu-22.04]
@@ -289,7 +297,7 @@ jobs:
         kokkos_version: ${{ fromJson(needs.build_and_cache_Kokkos.outputs.kokkos_version) }}
     timeout-minutes: 60
     name: C++ tests (multiple backends)
-    runs-on: ${{ matrix.os }}
+    runs-on: ${{ needs.determine_runner.outputs.runner_group }}
 
     steps:
       - uses: actions/setup-python@v5
diff --git a/.github/workflows/tests_linux_python.yml b/.github/workflows/tests_linux_python.yml
index 4a61f7a70d..dbb6ee1037 100644
--- a/.github/workflows/tests_linux_python.yml
+++ b/.github/workflows/tests_linux_python.yml
@@ -31,15 +31,21 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
+  determine_runner:
+    name: Determine runner type to use
+    uses: ./.github/workflows/determine-workflow-runner.yml
+    with:
+      default_runner: ubuntu-22.04
+
 
   pythontests:
+    needs: [determine_runner]
     strategy:
       matrix:
-        os: [ubuntu-22.04]
         pl_backend: ["lightning_qubit"]
     timeout-minutes: 75
     name: Python tests
-    runs-on: ${{ matrix.os }}
+    runs-on: ${{ needs.determine_runner.outputs.runner_group }}
 
     steps:
       - uses: actions/setup-python@v4
@@ -140,13 +146,13 @@ jobs:
 
 
   pythontestswithOpenBLAS:
+    needs: [determine_runner]
     strategy:
       matrix:
-        os: [ubuntu-22.04]
         pl_backend: ["lightning_qubit"]
     timeout-minutes: 75
     name: Python tests with OpenBLAS
-    runs-on: ${{ matrix.os }}
+    runs-on: ${{ needs.determine_runner.outputs.runner_group }}
 
     steps:
       - name: Checkout PennyLane-Lightning
@@ -242,12 +248,14 @@ jobs:
 
   build_and_cache_Kokkos:
     name: "Build and cache Kokkos"
+    needs: [determine_runner]
     uses: ./.github/workflows/build_and_cache_Kokkos_linux.yml
     with:
       os: ubuntu-22.04
+      runs_on: ${{ needs.determine_runner.outputs.runner_group }}
 
   pythontestswithKokkos:
-    needs: [build_and_cache_Kokkos]
+    needs: [build_and_cache_Kokkos, determine_runner]
     strategy:
       matrix:
         os: [ubuntu-22.04]
@@ -259,7 +267,7 @@ jobs:
             exec_model: OPENMP
     timeout-minutes: 75
     name: Python tests with Kokkos
-    runs-on: ${{ matrix.os }}
+    runs-on: ${{ needs.determine_runner.outputs.runner_group }}
 
     steps:
       - name: Checkout PennyLane-Lightning
diff --git a/.github/workflows/tests_without_binary.yml b/.github/workflows/tests_without_binary.yml
index 9854d39526..5a2a0b0d0c 100644
--- a/.github/workflows/tests_without_binary.yml
+++ b/.github/workflows/tests_without_binary.yml
@@ -25,13 +25,19 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
+  determine_runner:
+    name: Determine runner type to use
+    uses: ./.github/workflows/determine-workflow-runner.yml
+    with:
+      default_runner: ubuntu-22.04
+
   pythontests:
+    needs: [determine_runner]
     timeout-minutes: 30
     name: Python tests
-    runs-on: ${{ matrix.os }}
+    runs-on: ${{ needs.determine_runner.outputs.runner_group }}
     strategy:
       matrix:
-        os: [ubuntu-22.04]
         pl_backend: ["lightning_qubit", "lightning_kokkos", "lightning_gpu"]
 
     steps:
diff --git a/.github/workflows/wheel_linux_x86_64.yml b/.github/workflows/wheel_linux_x86_64.yml
index 0b42cdf085..c081fc7250 100644
--- a/.github/workflows/wheel_linux_x86_64.yml
+++ b/.github/workflows/wheel_linux_x86_64.yml
@@ -32,17 +32,22 @@ jobs:
     with:
       event_name: ${{ github.event_name }}
 
+  determine_runner:
+    name: Determine runner type to use
+    uses: ./.github/workflows/determine-workflow-runner.yml
+    with:
+      default_runner: ubuntu-latest
+
   build_dependencies:
-    needs: [set_wheel_build_matrix]
+    needs: [set_wheel_build_matrix, determine_runner]
     strategy:
       matrix:
-        os: [ubuntu-latest]
         exec_model: ${{ fromJson(needs.set_wheel_build_matrix.outputs.exec_model) }}
         kokkos_version: ${{ fromJson(needs.set_wheel_build_matrix.outputs.kokkos_version) }}
         container_img: ["quay.io/pypa/manylinux2014_x86_64"]
     timeout-minutes: 30
     name: Kokkos core (${{ matrix.exec_model }})
-    runs-on: ${{ matrix.os }}
+    runs-on: ${{ needs.determine_runner.outputs.runner_group }}
     container: ${{ matrix.container_img }}
 
     steps:
@@ -90,11 +95,10 @@ jobs:
           cd -
 
   linux-wheels-x86-64:
-    needs: [set_wheel_build_matrix, build_dependencies]
+    needs: [set_wheel_build_matrix, build_dependencies, determine_runner]
     strategy:
       fail-fast: false
       matrix:
-        os: [ubuntu-latest]
         arch: [x86_64]
         pl_backend: ["lightning_kokkos", "lightning_qubit"]
         cibw_build: ${{ fromJson(needs.set_wheel_build_matrix.outputs.python_version) }}
@@ -102,8 +106,8 @@ jobs:
         kokkos_version: ${{ fromJson(needs.set_wheel_build_matrix.outputs.kokkos_version) }}
         container_img: ["quay.io/pypa/manylinux2014_x86_64"]
     timeout-minutes: 30
-    name: ${{ matrix.os }}::${{ matrix.arch }} - ${{ matrix.pl_backend }} (Python ${{ fromJson('{ "cp39-*":"3.9","cp310-*":"3.10","cp311-*":"3.11","cp312-*":"3.12" }')[matrix.cibw_build] }})
-    runs-on: ${{ matrix.os }}
+    name: Linux::${{ matrix.arch }} - ${{ matrix.pl_backend }} (Python ${{ fromJson('{ "cp39-*":"3.9","cp310-*":"3.10","cp311-*":"3.11","cp312-*":"3.12" }')[matrix.cibw_build] }})
+    runs-on: ${{ needs.determine_runner.outputs.runner_group }}
     container: ${{ matrix.container_img }}
 
     steps:
diff --git a/.github/workflows/wheel_linux_x86_64_cuda.yml b/.github/workflows/wheel_linux_x86_64_cuda.yml
index 422507891b..aa1671fd7c 100644
--- a/.github/workflows/wheel_linux_x86_64_cuda.yml
+++ b/.github/workflows/wheel_linux_x86_64_cuda.yml
@@ -32,20 +32,26 @@ jobs:
     with:
       event_name: ${{ github.event_name }}
 
+  determine_runner:
+    name: Determine runner type to use
+    uses: ./.github/workflows/determine-workflow-runner.yml
+    with:
+      default_runner: ubuntu-latest
+
+
   linux-wheels-x86-64:
-    needs: [set_wheel_build_matrix]
+    needs: [set_wheel_build_matrix, determine_runner]
     strategy:
       fail-fast: false
       matrix:
-        os: [ubuntu-latest]
         arch: [x86_64]
         pl_backend: ["lightning_gpu"]
         cuda_version: ["12"]
         cibw_build: ${{ fromJson(needs.set_wheel_build_matrix.outputs.python_version) }}
         container_img: ["quay.io/pypa/manylinux2014_x86_64"]
     timeout-minutes: 30
-    name: ${{ matrix.os }}::${{ matrix.arch }} - ${{ matrix.pl_backend }} CUDA ${{ matrix.cuda_version }} (Python ${{ fromJson('{ "cp39-*":"3.9","cp310-*":"3.10","cp311-*":"3.11","cp312-*":"3.12" }')[matrix.cibw_build] }})
-    runs-on: ${{ matrix.os }}
+    name: Linux::${{ matrix.arch }} - ${{ matrix.pl_backend }} CUDA ${{ matrix.cuda_version }} (Python ${{ fromJson('{ "cp39-*":"3.9","cp310-*":"3.10","cp311-*":"3.11","cp312-*":"3.12" }')[matrix.cibw_build] }})
+    runs-on: ${{ needs.determine_runner.outputs.runner_group }}
     container: ${{ matrix.container_img }}
 
     steps:
diff --git a/.github/workflows/wheel_noarch.yml b/.github/workflows/wheel_noarch.yml
index 3d189a3931..209bdcff73 100644
--- a/.github/workflows/wheel_noarch.yml
+++ b/.github/workflows/wheel_noarch.yml
@@ -19,18 +19,25 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
+  determine_runner:
+    name: Determine runner type to use
+    uses: ./.github/workflows/determine-workflow-runner.yml
+    with:
+      default_runner: ubuntu-latest
+
+
   build-pure-python-wheel:
     if: |
       github.event_name != 'pull_request' ||
       contains(github.event.pull_request.labels.*.name, 'ci:build_wheels') ||
       github.event_name == 'workflow_dispatch'
+    needs: [determine_runner]
     strategy:
       matrix:
-        os: [ubuntu-latest]
         pl_backend: ["lightning_gpu", "lightning_kokkos", "lightning_qubit"]
     timeout-minutes: 30
-    name: ${{ matrix.os }} - Pure Python wheels - ${{ matrix.pl_backend }} (Python 3.9)
-    runs-on: ${{ matrix.os }}
+    name: Linux - Pure Python wheels - ${{ matrix.pl_backend }} (Python 3.9)
+    runs-on: ${{ needs.determine_runner.outputs.runner_group }}
 
     steps:
       - name: Checkout PennyLane-Lightning
diff --git a/pennylane_lightning/core/_version.py b/pennylane_lightning/core/_version.py
index 0198b9a3ca..a581fbb6e3 100644
--- a/pennylane_lightning/core/_version.py
+++ b/pennylane_lightning/core/_version.py
@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.36.0-dev45"
+__version__ = "0.36.0-dev46"

From cc57f14d0607a93af662dea1d48854a5a1b5381f Mon Sep 17 00:00:00 2001
From: Amintor Dusko <87949283+AmintorDusko@users.noreply.github.com>
Date: Fri, 3 May 2024 11:29:16 -0400
Subject: [PATCH 4/4] Port PR #705 changes (#711)

* port PR #705 changes

* Auto update version from '0.36.0-dev45' to '0.36.0-dev46'

* Update pennylane_lightning/lightning_qubit/lightning_qubit.py

Co-authored-by: Ali Asadi <10773383+maliasadi@users.noreply.github.com>

* Trigger CIs

* fix CodeFactor fail

* fix CodeFactor fail

* fix CodeFactor fail

* Auto update version from '0.36.0-dev46' to '0.36.0-dev47'

---------

Co-authored-by: ringo-but-quantum <github-ringo-but-quantum@xanadu.ai>
Co-authored-by: Ali Asadi <10773383+maliasadi@users.noreply.github.com>
---
 .github/CHANGELOG.md                          |  3 +
 pennylane_lightning/core/_version.py          |  2 +-
 .../lightning_qubit/lightning_qubit.py        | 61 +++++++++++++++----
 tests/new_api/test_device.py                  | 23 +++++++
 4 files changed, 75 insertions(+), 14 deletions(-)

diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index 700adc8e5b..3dd8f6b58c 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -105,6 +105,9 @@
 
 ### Bug fixes
 
+* Lightning Qubit once again respects the wire order specified on device instantiation.
+  [(#705)](https://github.com/PennyLaneAI/pennylane-lightning/pull/705)
+
 * `dynamic_one_shot` was refactored to use `SampleMP` measurements as a way to return the mid-circuit measurement samples. `LightningQubit`'s `simulate` is modified accordingly.
   [(#694)](https://github.com/PennyLaneAI/pennylane/pull/694)
 
diff --git a/pennylane_lightning/core/_version.py b/pennylane_lightning/core/_version.py
index a581fbb6e3..66cb20c0e9 100644
--- a/pennylane_lightning/core/_version.py
+++ b/pennylane_lightning/core/_version.py
@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.36.0-dev46"
+__version__ = "0.36.0-dev47"
diff --git a/pennylane_lightning/lightning_qubit/lightning_qubit.py b/pennylane_lightning/lightning_qubit/lightning_qubit.py
index 8350804c18..6fe1f307ae 100644
--- a/pennylane_lightning/lightning_qubit/lightning_qubit.py
+++ b/pennylane_lightning/lightning_qubit/lightning_qubit.py
@@ -87,7 +87,7 @@ def simulate(circuit: QuantumScript, state: LightningStateVector, mcmc: dict = N
     return LightningMeasurements(final_state, **mcmc).measure_final_state(circuit)
 
 
-def jacobian(circuit: QuantumTape, state: LightningStateVector, batch_obs=False):
+def jacobian(circuit: QuantumTape, state: LightningStateVector, batch_obs=False, wire_map=None):
     """Compute the Jacobian for a single quantum script.
 
     Args:
@@ -96,17 +96,21 @@ def jacobian(circuit: QuantumTape, state: LightningStateVector, batch_obs=False)
         batch_obs (bool): Determine whether we process observables in parallel when
             computing the jacobian. This value is only relevant when the lightning
             qubit is built with OpenMP. Default is False.
+        wire_map (Optional[dict]): a map from wire labels to simulation indices
 
     Returns:
         TensorLike: The Jacobian of the quantum script
     """
-    circuit = circuit.map_to_standard_wires()
+    if wire_map is not None:
+        [circuit], _ = qml.map_wires(circuit, wire_map)
     state.reset_state()
     final_state = state.get_final_state(circuit)
     return LightningAdjointJacobian(final_state, batch_obs=batch_obs).calculate_jacobian(circuit)
 
 
-def simulate_and_jacobian(circuit: QuantumTape, state: LightningStateVector, batch_obs=False):
+def simulate_and_jacobian(
+    circuit: QuantumTape, state: LightningStateVector, batch_obs=False, wire_map=None
+):
     """Simulate a single quantum script and compute its Jacobian.
 
     Args:
@@ -115,20 +119,26 @@ def simulate_and_jacobian(circuit: QuantumTape, state: LightningStateVector, bat
         batch_obs (bool): Determine whether we process observables in parallel when
             computing the jacobian. This value is only relevant when the lightning
             qubit is built with OpenMP. Default is False.
+        wire_map (Optional[dict]): a map from wire labels to simulation indices
 
     Returns:
         Tuple[TensorLike]: The results of the simulation and the calculated Jacobian
 
     Note that this function can return measurements for non-commuting observables simultaneously.
     """
-    circuit = circuit.map_to_standard_wires()
+    if wire_map is not None:
+        [circuit], _ = qml.map_wires(circuit, wire_map)
     res = simulate(circuit, state)
     jac = LightningAdjointJacobian(state, batch_obs=batch_obs).calculate_jacobian(circuit)
     return res, jac
 
 
 def vjp(
-    circuit: QuantumTape, cotangents: Tuple[Number], state: LightningStateVector, batch_obs=False
+    circuit: QuantumTape,
+    cotangents: Tuple[Number],
+    state: LightningStateVector,
+    batch_obs=False,
+    wire_map=None,
 ):
     """Compute the Vector-Jacobian Product (VJP) for a single quantum script.
     Args:
@@ -141,10 +151,13 @@ def vjp(
         batch_obs (bool): Determine whether we process observables in parallel when
             computing the VJP. This value is only relevant when the lightning
             qubit is built with OpenMP.
+        wire_map (Optional[dict]): a map from wire labels to simulation indices
+
     Returns:
         TensorLike: The VJP of the quantum script
     """
-    circuit = circuit.map_to_standard_wires()
+    if wire_map is not None:
+        [circuit], _ = qml.map_wires(circuit, wire_map)
     state.reset_state()
     final_state = state.get_final_state(circuit)
     return LightningAdjointJacobian(final_state, batch_obs=batch_obs).calculate_vjp(
@@ -153,7 +166,11 @@ def vjp(
 
 
 def simulate_and_vjp(
-    circuit: QuantumTape, cotangents: Tuple[Number], state: LightningStateVector, batch_obs=False
+    circuit: QuantumTape,
+    cotangents: Tuple[Number],
+    state: LightningStateVector,
+    batch_obs=False,
+    wire_map=None,
 ):
     """Simulate a single quantum script and compute its Vector-Jacobian Product (VJP).
     Args:
@@ -166,11 +183,14 @@ def simulate_and_vjp(
         batch_obs (bool): Determine whether we process observables in parallel when
             computing the jacobian. This value is only relevant when the lightning
             qubit is built with OpenMP.
+        wire_map (Optional[dict]): a map from wire labels to simulation indices
+
     Returns:
         Tuple[TensorLike]: The results of the simulation and the calculated VJP
     Note that this function can return measurements for non-commuting observables simultaneously.
     """
-    circuit = circuit.map_to_standard_wires()
+    if wire_map is not None:
+        [circuit], _ = qml.map_wires(circuit, wire_map)
     res = simulate(circuit, state)
     _vjp = LightningAdjointJacobian(state, batch_obs=batch_obs).calculate_vjp(circuit, cotangents)
     return res, _vjp
@@ -413,6 +433,8 @@ class LightningQubit(Device):
             qubit is built with OpenMP.
     """
 
+    # pylint: disable=too-many-instance-attributes
+
     _device_options = ("rng", "c_dtype", "batch_obs", "mcmc", "kernel_name", "num_burnin")
     _CPP_BINARY_AVAILABLE = LQ_CPP_BINARY_AVAILABLE
     _new_API = True
@@ -449,6 +471,11 @@ def __init__(  # pylint: disable=too-many-arguments
 
         super().__init__(wires=wires, shots=shots)
 
+        if isinstance(wires, int):
+            self._wire_map = None  # should just use wires as is
+        else:
+            self._wire_map = {w: i for i, w in enumerate(self.wires)}
+
         self._statevector = LightningStateVector(num_wires=len(self.wires), dtype=c_dtype)
 
         # TODO: Investigate usefulness of creating numpy random generator
@@ -568,7 +595,8 @@ def execute(
         }
         results = []
         for circuit in circuits:
-            circuit = circuit.map_to_standard_wires()
+            if self._wire_map is not None:
+                [circuit], _ = qml.map_wires(circuit, self._wire_map)
             results.append(simulate(circuit, self._statevector, mcmc=mcmc))
 
         return tuple(results)
@@ -613,8 +641,10 @@ def compute_derivatives(
             Tuple: The jacobian for each trainable parameter
         """
         batch_obs = execution_config.device_options.get("batch_obs", self._batch_obs)
+
         return tuple(
-            jacobian(circuit, self._statevector, batch_obs=batch_obs) for circuit in circuits
+            jacobian(circuit, self._statevector, batch_obs=batch_obs, wire_map=self._wire_map)
+            for circuit in circuits
         )
 
     def execute_and_compute_derivatives(
@@ -633,7 +663,10 @@ def execute_and_compute_derivatives(
         """
         batch_obs = execution_config.device_options.get("batch_obs", self._batch_obs)
         results = tuple(
-            simulate_and_jacobian(c, self._statevector, batch_obs=batch_obs) for c in circuits
+            simulate_and_jacobian(
+                c, self._statevector, batch_obs=batch_obs, wire_map=self._wire_map
+            )
+            for c in circuits
         )
         return tuple(zip(*results))
 
@@ -686,7 +719,7 @@ def compute_vjp(
         """
         batch_obs = execution_config.device_options.get("batch_obs", self._batch_obs)
         return tuple(
-            vjp(circuit, cots, self._statevector, batch_obs=batch_obs)
+            vjp(circuit, cots, self._statevector, batch_obs=batch_obs, wire_map=self._wire_map)
             for circuit, cots in zip(circuits, cotangents)
         )
 
@@ -708,7 +741,9 @@ def execute_and_compute_vjp(
         """
         batch_obs = execution_config.device_options.get("batch_obs", self._batch_obs)
         results = tuple(
-            simulate_and_vjp(circuit, cots, self._statevector, batch_obs=batch_obs)
+            simulate_and_vjp(
+                circuit, cots, self._statevector, batch_obs=batch_obs, wire_map=self._wire_map
+            )
             for circuit, cots in zip(circuits, cotangents)
         )
         return tuple(zip(*results))
diff --git a/tests/new_api/test_device.py b/tests/new_api/test_device.py
index bf27cbcb47..cbbfb286e5 100644
--- a/tests/new_api/test_device.py
+++ b/tests/new_api/test_device.py
@@ -456,6 +456,29 @@ def test_custom_wires(self, phi, theta, wires):
         assert np.allclose(result[0], np.cos(phi))
         assert np.allclose(result[1], np.cos(phi) * np.cos(theta))
 
+    @pytest.mark.parametrize(
+        "wires, wire_order", [(3, (0, 1, 2)), (("a", "b", "c"), ("a", "b", "c"))]
+    )
+    def test_probs_different_wire_orders(self, wires, wire_order):
+        """Test that measuring probabilities works with custom wires."""
+
+        dev = LightningDevice(wires=wires)
+
+        op = qml.Hadamard(wire_order[1])
+
+        tape = QuantumScript([op], [qml.probs(wires=(wire_order[0], wire_order[1]))])
+
+        res = dev.execute(tape)
+        assert qml.math.allclose(res, np.array([0.5, 0.5, 0.0, 0.0]))
+
+        tape2 = QuantumScript([op], [qml.probs(wires=(wire_order[1], wire_order[2]))])
+        res2 = dev.execute(tape2)
+        assert qml.math.allclose(res2, np.array([0.5, 0.0, 0.5, 0.0]))
+
+        tape3 = QuantumScript([op], [qml.probs(wires=(wire_order[1], wire_order[0]))])
+        res3 = dev.execute(tape3)
+        assert qml.math.allclose(res3, np.array([0.5, 0.0, 0.5, 0.0]))
+
 
 @pytest.mark.parametrize("batch_obs", [True, False])
 class TestDerivatives: