Add stateprep support to lightning.tensor - MPS

.github/CHANGELOG.md

@@ -2,6 +2,9 @@
 
 ### New features since last release
 
+* Add `qml.StatePrep()` and `qml.QubitStateVector()` support to `lightning.tensor`.
+  [(#849)](https://github.com/PennyLaneAI/pennylane-lightning/pull/849)
+
 * Add LightningGPU Linux (aarch64+GraceHopper) wheels to PyPI.
   [(#815)](https://github.com/PennyLaneAI/pennylane-lightning/pull/815)
 

pennylane_lightning/core/src/simulators/lightning_tensor/tncuda/MPSTNCuda.hpp

@@ -64,7 +64,7 @@ class MPSTNCuda final : public TNCudaBase<Precision, MPSTNCuda<Precision>> {
     MPSStatus MPSInitialized_ = MPSStatus::MPSInitNotSet;
 
     const std::size_t maxBondDim_;
-
+    const std::vector<std::size_t> bondDims_;
     const std::vector<std::vector<std::size_t>> sitesModes_;
     const std::vector<std::vector<std::size_t>> sitesExtents_;
     const std::vector<std::vector<int64_t>> sitesExtents_int64_;
@@ -86,19 +86,25 @@ class MPSTNCuda final : public TNCudaBase<Precision, MPSTNCuda<Precision>> {
     explicit MPSTNCuda(const std::size_t numQubits,
                        const std::size_t maxBondDim)
         : BaseType(numQubits), maxBondDim_(maxBondDim),
-          sitesModes_(setSitesModes_()), sitesExtents_(setSitesExtents_()),
+          bondDims_(setBondDims_()), sitesModes_(setSitesModes_()),
+          sitesExtents_(setSitesExtents_()),
           sitesExtents_int64_(setSitesExtents_int64_()) {
         initTensors_();
+        reset();
+        appendInitialMPSState_();
     }
 
     // TODO: Add method to the constructor to allow users to select methods at
     // runtime in the C++ layer
     explicit MPSTNCuda(const std::size_t numQubits,
                        const std::size_t maxBondDim, DevTag<int> dev_tag)
         : BaseType(numQubits, dev_tag), maxBondDim_(maxBondDim),
-          sitesModes_(setSitesModes_()), sitesExtents_(setSitesExtents_()),
+          bondDims_(setBondDims_()), sitesModes_(setSitesModes_()),
+          sitesExtents_(setSitesExtents_()),
           sitesExtents_int64_(setSitesExtents_int64_()) {
         initTensors_();
+        reset();
+        appendInitialMPSState_();
     }
 
     ~MPSTNCuda() = default;
@@ -162,6 +168,31 @@ class MPSTNCuda final : public TNCudaBase<Precision, MPSTNCuda<Precision>> {
         setBasisState(zeroState);
     }
 
+    /**
+     * @brief Update the ith MPS site data.
+     *
+     * @param site_idx Index of the MPS site.
+     * @param host_data Pointer to the data on host.
+     * @param host_data_size Length of the data.
+     */
+    void updateIthMPSSiteData(const std::size_t site_idx,
+                              const ComplexT *host_data,
+                              std::size_t host_data_size) {
+        PL_ABORT_IF_NOT(
+            site_idx < BaseType::getNumQubits(),
+            "The site index should be less than the number of qubits.");
+
+        const std::size_t idx = BaseType::getNumQubits() - site_idx - 1;
+        PL_ABORT_IF_NOT(
+            host_data_size == tensors_[idx].getDataBuffer().getLength(),
+            "The length of the host data should match its copy on the device.");
+
+        tensors_[idx].getDataBuffer().zeroInit();
+
+        tensors_[idx].getDataBuffer().CopyHostDataToGpu(host_data,
+                                                        host_data_size);
+    }
+
     /**
      * @brief Update quantum state with a basis state.
      * NOTE: This API assumes the bond vector is a standard basis vector
@@ -192,18 +223,13 @@ class MPSTNCuda final : public TNCudaBase<Precision, MPSTNCuda<Precision>> {
             if (i == 0) {
                 target = basisState[idx];
             } else {
-                target = basisState[idx] == 0 ? 0 : maxBondDim_;
+                target = basisState[idx] == 0 ? 0 : bondDims_[i - 1];
             }
 
             PL_CUDA_IS_SUCCESS(
                 cudaMemcpy(&tensors_[i].getDataBuffer().getData()[target],
                            &value_cu, sizeof(CFP_t), cudaMemcpyHostToDevice));
         }
-
-        if (MPSInitialized_ == MPSStatus::MPSInitNotSet) {
-            MPSInitialized_ = MPSStatus::MPSInitSet;
-            updateQuantumStateMPS_();
-        }
     };
 
     /**
@@ -307,6 +333,27 @@ class MPSTNCuda final : public TNCudaBase<Precision, MPSTNCuda<Precision>> {
     }
 
   private:
+    /**
+     * @brief Return bondDims to the member initializer
+     * NOTE: This method only works for the open boundary condition
+     * @return std::vector<std::size_t>
+     */
+    std::vector<std::size_t> setBondDims_() {
+        std::vector<std::size_t> localBondDims(BaseType::getNumQubits() - 1,
+                                               maxBondDim_);
+
+        const std::size_t ubDim = log2(maxBondDim_);
+        for (std::size_t i = 0; i < localBondDims.size(); i++) {
+            const std::size_t bondDim =
+                std::min(i + 1, BaseType::getNumQubits() - i - 1);
+
+            if (bondDim <= ubDim) {
+                localBondDims[i] = std::size_t{1} << bondDim;
+            }
+        }
+        return localBondDims;
+    }
+
     /**
      * @brief Return siteModes to the member initializer
      * NOTE: This method only works for the open boundary condition
@@ -348,15 +395,16 @@ class MPSTNCuda final : public TNCudaBase<Precision, MPSTNCuda<Precision>> {
             if (i == 0) {
                 // Leftmost site (state mode, shared mode)
                 localSiteExtents = std::vector<std::size_t>(
-                    {BaseType::getQubitDims()[i], maxBondDim_});
+                    {BaseType::getQubitDims()[i], bondDims_[i]});
             } else if (i == BaseType::getNumQubits() - 1) {
                 // Rightmost site (shared mode, state mode)
                 localSiteExtents = std::vector<std::size_t>(
-                    {maxBondDim_, BaseType::getQubitDims()[i]});
+                    {bondDims_[i - 1], BaseType::getQubitDims()[i]});
             } else {
                 // Interior sites (state mode, state mode, shared mode)
                 localSiteExtents = std::vector<std::size_t>(
-                    {maxBondDim_, BaseType::getQubitDims()[i], maxBondDim_});
+                    {bondDims_[i - 1], BaseType::getQubitDims()[i],
+                     bondDims_[i]});
             }
             localSitesExtents.push_back(std::move(localSiteExtents));
         }
@@ -394,11 +442,11 @@ class MPSTNCuda final : public TNCudaBase<Precision, MPSTNCuda<Precision>> {
     }
 
     /**
-     * @brief Update quantumState (cutensornetState_t) with data provided by a
-     * user
+     * @brief Append initial MPS sites to the compute graph with data provided
+     * by a user
      *
      */
-    void updateQuantumStateMPS_() {
+    void appendInitialMPSState_() {
         PL_CUTENSORNET_IS_SUCCESS(cutensornetStateInitializeMPS(
             /*const cutensornetHandle_t */ BaseType::getTNCudaHandle(),
             /*cutensornetState_t*/ BaseType::getQuantumState(),

pennylane_lightning/core/src/simulators/lightning_tensor/tncuda/bindings/LTensorTNCudaBindings.hpp

@@ -74,6 +74,18 @@ void registerBackendClassSpecificBindings(PyClass &pyclass) {
                 tensor_network.getData(data_ptr, state.size());
             },
             "Copy StateVector data into a Numpy array.")
+        .def(
+            "updateMPSSitesData",
+            [](TensorNet &tensor_network, std::vector<np_arr_c> &tensors) {
+                for (std::size_t idx = 0; idx < tensors.size(); idx++) {
+                    py::buffer_info numpyArrayInfo = tensors[idx].request();
+                    auto *data_ptr = static_cast<std::complex<PrecisionT> *>(
+                        numpyArrayInfo.ptr);
+                    tensor_network.updateIthMPSSiteData(idx, data_ptr,
+                                                        tensors[idx].size());
+                }
+            },
+            "Pass MPS site data to the C++ backend")
         .def(
             "setBasisState",
             [](TensorNet &tensor_network,

pennylane_lightning/core/src/simulators/lightning_tensor/tncuda/measurements/tests/Test_MPSTNCuda_Expval.cpp

@@ -203,7 +203,7 @@ TEMPLATE_TEST_CASE("[PauliZ]", "[MPSTNCuda_Expval]", float, double) {
             PrecisionT ref = -0.2115276040475712;
 
             REQUIRE_THAT(res, Catch::Matchers::WithinRel(
-                                  ref, static_cast<PrecisionT>(cutoff)));
+                                  ref, static_cast<PrecisionT>(0.1)));
         }
     }
 }

pennylane_lightning/core/src/simulators/lightning_tensor/tncuda/tests/Tests_MPSTNCuda.cpp

@@ -58,6 +58,67 @@ TEMPLATE_PRODUCT_TEST_CASE("MPSTNCuda::Constructibility",
     }
 }
 
+TEMPLATE_TEST_CASE("MPSTNCuda::setIthMPSSite", "[MPSTNCuda]", float, double) {
+    SECTION("Set MPS site with wrong site index") {
+        const std::size_t num_qubits = 3;
+        const std::size_t maxBondDim = 3;
+        const std::size_t siteIdx = 3;
+
+        MPSTNCuda<TestType> mps_state{num_qubits, maxBondDim};
+
+        std::vector<std::complex<TestType>> site_data(1, {0.0, 0.0});
+
+        REQUIRE_THROWS_WITH(
+            mps_state.updateIthMPSSiteData(siteIdx, site_data.data(),
+                                           site_data.size()),
+            Catch::Matchers::Contains(
+                "The site index should be less than the number of qubits."));
+    }
+
+    SECTION("Set MPS site with wrong site data size") {
+        const std::size_t num_qubits = 3;
+        const std::size_t maxBondDim = 3;
+        const std::size_t siteIdx = 0;
+
+        MPSTNCuda<TestType> mps_state{num_qubits, maxBondDim};
+
+        std::vector<std::complex<TestType>> site_data(1, {0.0, 0.0});
+
+        REQUIRE_THROWS_WITH(
+            mps_state.updateIthMPSSiteData(siteIdx, site_data.data(),
+                                           site_data.size()),
+            Catch::Matchers::Contains("The length of the host data should "
+                                      "match its copy on the device."));
+    }
+
+    SECTION("Set MPS sites") {
+        const std::size_t num_qubits = 2;
+        const std::size_t maxBondDim = 3;
+
+        MPSTNCuda<TestType> mps_state{num_qubits, maxBondDim};
+
+        mps_state.reset(); // Reset the state to zero state
+
+        std::vector<std::complex<TestType>> site0_data(4, {0.0, 0.0}); // MSB
+        std::vector<std::complex<TestType>> site1_data(4, {0.0, 0.0}); // LSB
+
+        site0_data[2] = {1.0, 0.0};
+        site1_data[1] = {1.0, 0.0};
+
+        mps_state.updateIthMPSSiteData(0, site0_data.data(), site0_data.size());
+        mps_state.updateIthMPSSiteData(1, site1_data.data(), site1_data.size());
+
+        auto results = mps_state.getDataVector();
+
+        std::vector<std::complex<TestType>> expected_state(
+            std::size_t{1} << num_qubits, std::complex<TestType>({0.0, 0.0}));
+
+        expected_state[3] = {1.0, 0.0};
+
+        CHECK(expected_state == Pennylane::Util::approx(results));
+    }
+}
+
 TEMPLATE_TEST_CASE("MPSTNCuda::SetBasisStates() & reset()", "[MPSTNCuda]",
                    float, double) {
     std::vector<std::vector<std::size_t>> basisStates = {

pennylane_lightning/lightning_tensor/_tensornet.py

@@ -21,12 +21,52 @@
 except ImportError:
     pass
 
+from itertools import product
 
 import numpy as np
 import pennylane as qml
 from pennylane import BasisState, DeviceError, StatePrep
 from pennylane.ops.op_math import Adjoint
 from pennylane.tape import QuantumScript
+from pennylane.wires import Wires
+
+
+def svd_split(M, bond_dim):
+    """SVD split a matrix into a matrix product state via numpy linalg. Note that this function is to be moved to the C++ layer."""
+    U, S, Vd = np.linalg.svd(M, full_matrices=False)
+    U = U @ np.diag(S)  # Append singular values to U
+    bonds = len(S)
+    Vd = Vd.reshape(bonds, 2, -1)
+    U = U.reshape((-1, 2, bonds))
+
+    # keep only chi bonds
+    chi = np.min([bonds, bond_dim])
+    U, S, Vd = U[:, :, :chi], S[:chi], Vd[:chi]
+    return U, Vd
+
+
+def dense_to_mps(psi, n_wires, bond_dim):
+    """Convert a dense state vector to a matrix product state."""
+    Ms = []
+
+    psi = np.reshape(psi, (2, -1))  # split psi[2, 2, 2, 2..] = psi[2, (2x2x2...)]
+    U, Vd = svd_split(psi, bond_dim)  # psi[2, (2x2x..)] = U[2, mu] Vd[mu, (2x2x2x..)]
+
+    Ms.append(U)
+    bondL = Vd.shape[0]
+    psi = Vd
+
+    for _ in range(n_wires - 2):
+        psi = np.reshape(psi, (2 * bondL, -1))  # reshape psi[2 * bondL, (2x2x2...)]
+        U, Vd = svd_split(psi, bond_dim)  # psi[2, (2x2x..)] = U[2, mu] Vd[mu, (2x2x2x..)]
+        Ms.append(U)
+
+        psi = Vd
+        bondL = Vd.shape[0]
+
+    Ms.append(Vd)
+
+    return Ms
 
 
 # pylint: disable=too-many-instance-attributes
@@ -112,6 +152,54 @@ def reset_state(self):
         # init the quantum state to |00..0>
         self._tensornet.reset()
 
+    def _preprocess_state_vector(self, state, device_wires):
+        """Convert a specified state to a full internal state vector.
+
+        Args:
+            state (array[complex]): normalized input state of length ``2**len(wires)``
+            device_wires (Wires): wires that get initialized in the state
+
+        Returns:
+            array[complex]: normalized input state of length ``2**len(wires)``
+        """
+        output_shape = [2] * self._num_wires
+        # special case for integral types
+        if state.dtype.kind == "i":
+            state = np.array(state, dtype=self.dtype)
+
+        if len(device_wires) == self._num_wires and Wires(sorted(device_wires)) == device_wires:
+            return np.reshape(state, output_shape).ravel(order="C")
+
+        # generate basis states on subset of qubits via the cartesian product
+        basis_states = np.array(list(product([0, 1], repeat=len(device_wires))))
+
+        # get basis states to alter on full set of qubits
+        unravelled_indices = np.zeros((2 ** len(device_wires), self._num_wires), dtype=int)
+        unravelled_indices[:, device_wires] = basis_states
+
+        # get indices for which the state is changed to input state vector elements
+        ravelled_indices = np.ravel_multi_index(unravelled_indices.T, [2] * self._num_wires)
+
+        # get full state vector to be factorized into MPS
+        full_state = np.zeros(2**self._num_wires, dtype=self.dtype)
+        for i, value in enumerate(state):
+            full_state[ravelled_indices[i]] = value
+        return np.reshape(full_state, output_shape).ravel(order="C")
+
+    def _apply_state_vector(self, state, device_wires: Wires):
+        """Convert a specified state to MPS sites.
+        Args:
+            state (array[complex]): normalized input state of length ``2**len(wires)``
+                or broadcasted state of shape ``(batch_size, 2**len(wires))``
+            device_wires (Wires): wires that get initialized in the state
+        """
+
+        state = self._preprocess_state_vector(state, device_wires)
+
+        M = dense_to_mps(state, self._num_wires, self._max_bond_dim)
+
+        self._tensornet.updateMPSSitesData(M)
+
     def _apply_basis_state(self, state, wires):
         """Initialize the quantum state in a specified computational basis state.
 
@@ -177,10 +265,9 @@ def apply_operations(self, operations):
         # State preparation is currently done in Python
         if operations:  # make sure operations[0] exists
             if isinstance(operations[0], StatePrep):
-                raise DeviceError(
-                    "lightning.tensor does not support initialization with a state vector."
-                )
-            if isinstance(operations[0], BasisState):
+                self._apply_state_vector(operations[0].parameters[0].copy(), operations[0].wires)
+                operations = operations[1:]
+            elif isinstance(operations[0], BasisState):
                 self._apply_basis_state(operations[0].parameters[0], operations[0].wires)
                 operations = operations[1:]