Optimize memory peak for _preprocess_state_vector in LightningTensor

pennylane_lightning/core/_version.py

@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.39.0-dev42"
+__version__ = "0.39.0-dev43"

pennylane_lightning/lightning_tensor/_tensornet.py

@@ -21,8 +21,6 @@
 except ImportError:
     pass
 
-from itertools import product
-
 import numpy as np
 import pennylane as qml
 from pennylane import BasisState, DeviceError, StatePrep
@@ -223,20 +221,35 @@
         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))))
+        local_dev_wires = np.array(device_wires.tolist())[::-1]
+
+        # generate basis states on subset of qubits via broadcasting
+        base = np.tile([0, 1], 2 ** (len(local_dev_wires) - 1)).astype(dtype=np.int64)
+        indexes = np.zeros(2 ** (len(local_dev_wires)), dtype=np.int64)
+
+        max_dev_wire = self._num_wires - 1
 
         # 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
+        for i, wire in enumerate(local_dev_wires):
+
+            # get indices for which the state is changed to input state vector elements
+            indexes += base * 2 ** (max_dev_wire - wire)
+
+            if i == len(local_dev_wires) - 1:
+                continue
+
+            two_n = 2 ** (i + 1)
+            base = base.reshape(-1, two_n * 2)
+            swaper_A = two_n // 2
+            swaper_B = swaper_A + two_n
 
-        # 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)
+            base[:, swaper_A:swaper_B] = base[:, swaper_A:swaper_B][:, ::-1]
+            base = base.reshape(-1)
 
         # 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
+            full_state[indexes[i]] = value
         return np.reshape(full_state, output_shape).ravel(order="C")
 
     def _apply_state_vector(self, state, device_wires: Wires):
@@ -285,7 +298,7 @@
             None
         """
         # TODO: Discuss if public interface for max_mpo_bond_dim argument
-        max_mpo_bond_dim = 2 ** len(wires)  # Exact SVD decomposition for MPO
+        max_mpo_bond_dim = self._max_bond_dim
 
         # Get sorted wires and MPO site tensor
         mpos, sorted_wires = gate_matrix_decompose(