Remove mps limit custom mps

00_Playground/LT-mps_custom_bond_dim.py

@@ -0,0 +1,131 @@
+import numpy as np
+import pennylane as qml
+
+
+# Function to define the bond dim for MPS on base to the number of qubits.
+def setBondDims(numQubits, maxBondDim):
+
+    log_maxBondDim = np.log2(maxBondDim)
+    limit_dimension = 2 ** int(log_maxBondDim)
+    localBondDims = np.ones(numQubits -1) * limit_dimension
+
+    for i, val in enumerate(localBondDims):
+        bondDim = min(i+1, numQubits - i - 1)
+        if bondDim <= log_maxBondDim:
+            localBondDims[i] = 2**bondDim
+
+    return localBondDims
+
+def setSitesModes(numQubits):
+    localSitesModes = []
+
+    for i in range(numQubits):
+        if i == 0:
+            localSite = [i, i+numQubits, -1]
+        elif i == numQubits - 1:
+            localSite = [i+numQubits-1, i, -1]
+        else:
+            localSite = [i + numQubits - 1, i , i + numQubits]
+
+        localSitesModes.append(localSite)
+
+    localSitesModes = np.array(localSitesModes)
+
+    return localSitesModes
+
+def setSitesExtents(numQubits, bondDims):
+    qubitDims = np.ones(numQubits,dtype=int) * 2
+    localSiteExtents = []
+    for i in range(numQubits):
+        if i == 0:
+            localSite = [qubitDims[i], bondDims[i], -1]
+        elif i == numQubits - 1:
+            localSite = [bondDims[i-1], qubitDims[i], -1]
+        else:
+            localSite = [bondDims[i-1], qubitDims[i], bondDims[i]]
+
+        localSiteExtents.append(localSite)
+
+    localSiteExtents = np.array(localSiteExtents).astype(int)
+
+    return localSiteExtents
+
+# Function to print an MPS.
+def print_mps(mps, full=False, name=None):
+    """Print the MPOs."""
+
+    max_sites_length = 10
+
+    print("-"*100)
+    if name is not None:
+        print("MPS name:", name)
+
+    max_bond_dim = 0
+    for i, site in enumerate(mps):
+        max_bond_dim = max(max_bond_dim, site.shape[-1])
+
+    print("Sites:", len(mps), "| Max bond dim:", max_bond_dim)
+
+    max_rows = len(mps)//max_sites_length
+    if len(mps) % max_sites_length != 0:
+        max_rows += 1
+
+    for rows in range(max_rows):
+        for i in range(rows*max_sites_length, (rows+1)*max_sites_length):
+            print(f"| {'site '+str(i):^12}",end="")
+
+        print()
+        for i, site in enumerate(mps[rows*max_sites_length:(rows+1)*max_sites_length]):
+            print(f"| {str(site.shape):^12}", end="")
+        print()
+        print()
+
+    if full:
+        print("~"*100)
+        for i, site in enumerate(mps):
+            # print("Site", i, ":\n", site)
+            print("Site", i, " | Shape:", site.shape)            
+            print(site)
+            print("~"*100)
+    print("-"*100)
+
+# Create a random MPS with the custom dimensions
+def create_MPS_with_custom_bondDims(numQubits,bondDims):
+
+    sitesExtents = setSitesExtents(numQubits, bondDims)
+
+    MPS_example = []
+
+    for T_shape in sitesExtents:
+
+        if T_shape[-1] == -1:
+            T_shape = T_shape[:-1]
+
+        MPS_site = np.random.rand(*(tuple(T_shape)))
+
+        MPS_example.append(MPS_site)
+
+    return MPS_example
+
+if __name__ == "__main__":
+
+    wires = 4
+    bond_dims = [5,7,5]
+
+    MPS_example = create_MPS_with_custom_bondDims(wires, bond_dims)
+
+    method = "mps"
+
+    dev = qml.device("lightning.tensor", wires=wires, method=method, cutoff=1e-7, bond_dim=bond_dims)
+
+    dev_wires = dev.wires.tolist()
+
+    @qml.qnode(dev)
+    def circuit(custom_MPS=MPS_example):
+
+        qml.MPSPrep(mps=custom_MPS, wires=dev_wires)
+
+        return qml.expval(qml.PauliZ(0))
+
+    print("result",circuit())
+

pennylane_lightning/core/_version.py

@@ -16,4 +16,4 @@
 Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.41.0-dev20"
+__version__ = "0.41.0-dev21"

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

@@ -80,6 +80,11 @@ class MPSTNCuda final : public TNCuda<Precision, MPSTNCuda<Precision>> {
                        const std::size_t maxBondDim)
         : BaseType(numQubits, maxBondDim) {}
 
+    explicit MPSTNCuda(const std::size_t numQubits,
+                       const std::size_t maxBondDim,
+                       const std::vector<std::size_t> &bondDims)
+        : BaseType(numQubits, maxBondDim, bondDims) {}
+
     explicit MPSTNCuda(const std::size_t numQubits,
                        const std::size_t maxBondDim, DevTag<int> dev_tag)
         : BaseType(numQubits, dev_tag, maxBondDim) {}

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

@@ -79,6 +79,23 @@ class TNCuda : public TNCudaBase<PrecisionT, Derived> {
                           // of cutensornet.
     }
 
+    // Constructor with bondDims for custom MPS
+    explicit TNCuda(std::size_t numQubits, std::size_t maxBondDim,
+                    const std::vector<std::size_t> &bondDims)
+        : BaseType(numQubits), maxBondDim_(maxBondDim), bondDims_(bondDims),
+          sitesModes_(setSitesModes_()), sitesExtents_(setSitesExtents_()),
+          sitesExtents_int64_(setSitesExtents_int64_()),
+          cublascaller_(make_shared_cublas_caller()),
+          gate_cache_(std::make_shared<TNCudaGateCache<PrecisionT>>(
+              BaseType::getDevTag())) {
+        initTensors_();
+        appendInitialMPSState_(
+            getSitesExtentsPtr()
+                .data()); // This API works for Exact Tensor Network as well,
+                          // given sitesExtents_ settings meet the requirement
+                          // of cutensornet.
+    }
+
     explicit TNCuda(const std::size_t numQubits, DevTag<int> dev_tag,
                     const std::size_t maxBondDim = 1)
         : BaseType(numQubits, dev_tag.getDeviceID(), dev_tag.getStreamID()),

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

@@ -121,6 +121,9 @@ void registerBackendClassSpecificBindingsMPS(PyClass &pyclass) {
 
     pyclass
         .def(py::init<std::size_t, std::size_t>()) // num_qubits, max_bond_dim
+        .def(py::init<std::size_t, std::size_t,
+                      std::vector<std::size_t>>()) // num_qubits, max_bond_dim,
+                                                   // list_bond_dim
         .def(py::init<std::size_t, std::size_t,
                       DevTag<int>>()) // num_qubits, max_bond_dim, dev-tag
         .def(
@@ -150,6 +153,25 @@ void registerBackendClassSpecificBindingsMPS(PyClass &pyclass) {
                 }
 
                 const auto &MPS_shape_dest = tensor_network.getSitesExtents();
+                // print MPS_shape_dest and MPS_shape_source
+
+                std::cout << "MPS_shape_dest: " << std::endl;
+                for (const auto &shape : MPS_shape_dest) {
+                    for (const auto &dim : shape) {
+                        std::cout << dim << " ";
+                    }
+                    std::cout << std::endl;
+                }
+
+                std::cout << "MPS_shape_source: " << std::endl;
+                for (const auto &shape : MPS_shape_source) {
+                    for (const auto &dim : shape) {
+                        std::cout << dim << " ";
+                    }
+                    std::cout << std::endl;
+                }
+
+
                 MPSShapeCheck(MPS_shape_dest, MPS_shape_source);
 
                 for (std::size_t idx = 0; idx < tensors.size(); idx++) {

pennylane_lightning/lightning_tensor/_tensornet.py

@@ -175,7 +175,12 @@
             self._max_bond_dim = kwargs.get("max_bond_dim", 128)
             self._cutoff = kwargs.get("cutoff", 0)
             self._cutoff_mode = kwargs.get("cutoff_mode", "abs")
-            self._tensornet = self._tensornet_dtype()(self._num_wires, self._max_bond_dim)
+            self._bond_dim = kwargs.get("bond_dim", None)
+
+            if self._bond_dim is not None:
+                self._tensornet = self._tensornet_dtype()(self._num_wires, self._max_bond_dim, self._bond_dim)
+            else:
+                self._tensornet = self._tensornet_dtype()(self._num_wires, self._max_bond_dim)
         elif self._method == "tn":
             self._tensornet = self._tensornet_dtype()(self._num_wires)
         else:

pennylane_lightning/lightning_tensor/lightning_tensor.py

@@ -287,7 +287,7 @@ def circuit(num_qubits):
     # pylint: disable=too-many-instance-attributes
 
     _device_options = {
-        "mps": ("backend", "max_bond_dim", "cutoff", "cutoff_mode"),
+        "mps": ("backend", "max_bond_dim", "cutoff", "cutoff_mode", "bond_dim"),
         "tn": ("backend"),
     }
 
@@ -340,12 +340,15 @@ def __init__(
                     f"Unexpected argument: {arg} during initialization of the lightning.tensor device."
                 )
 
+        self._custom_MPS = False
+
         if not accepted_backends(self._backend):
             raise ValueError(f"Unsupported backend: {self._backend}")
         if self._method == "mps":
             self._max_bond_dim = kwargs.get("max_bond_dim", 128)
             self._cutoff = kwargs.get("cutoff", 0)
             self._cutoff_mode = kwargs.get("cutoff_mode", "abs")
+            self._bond_dim = kwargs.get("bond_dim", None)
 
             if not isinstance(self._max_bond_dim, int) or self._max_bond_dim < 1:
                 raise ValueError("The maximum bond dimension must be an integer greater than 0.")
@@ -354,6 +357,13 @@ def __init__(
             if self._cutoff_mode not in ["rel", "abs"]:
                 raise ValueError(f"Unsupported cutoff mode: {self._cutoff_mode}")
 
+            if self._bond_dim is not None:
+                self._custom_MPS = True
+                if not isinstance(self._bond_dim, list):
+                    raise ValueError("The bond dimension must be specified as a list.")
+                if len(self._bond_dim) != self._num_wires - 1:
+                    raise ValueError("The bond dimension list must have length equal to the number of wires - 1.")
+
     @property
     def name(self):
         """The name of the device."""
@@ -382,6 +392,17 @@ def c_dtype(self):
     def _tensornet(self, num_wires):
         """Return the tensornet object."""
         if self.method == "mps":
+            if self._custom_MPS:
+                return LightningTensorNet(
+                    num_wires,
+                    self._method,
+                    self._c_dtype,
+                    device_name=self.name,
+                    max_bond_dim=self._max_bond_dim,
+                    cutoff=self._cutoff,
+                    cutoff_mode=self._cutoff_mode,
+                    bond_dim=self._bond_dim,
+                )
             return LightningTensorNet(
                 num_wires,
                 self._method,