Add docstrings

src/qibotn/mps_contraction_helper.py

@@ -7,17 +7,8 @@ class MPSContractionHelper:
     """A helper class to compute various quantities for a given MPS.
 
     Interleaved format is used to construct the input args for `cuquantum.contract`.
-    A concrete example on how the modes are populated for a 7-site MPS is provided
-    below::
-
-              0     2     4     6     8    10     12    14
-        bra -----A-----B-----C-----D-----E-----F-----G-----
-                 |     |     |     |     |     |     |
-                1|    3|    5|    7|    9|   11|   13|
-                 |     |     |     |     |     |     |
-        ket -----a-----b-----c-----d-----e-----f-----g-----
-              15    16    17    18    19    20    21    22
     
+    Reference: https://github.com/NVIDIA/cuQuantum/blob/main/python/samples/cutensornet/tn_algorithms/mps_algorithms.ipynb
 
     The follwing compute quantities are supported:
 
@@ -26,10 +17,6 @@ class MPSContractionHelper:
         - the expectation value for a given operator.
         - the equivalent state vector after multiplying an MPO to an MPS.
 
-    Note that for the nth MPS tensor (rank-3), the modes of the tensor are expected to be `(i,p,j)`
-    where i denotes the bonding mode with the (n-1)th tensor, p denotes the physical mode for the qubit and
-    j denotes the bonding mode with the (n+1)th tensor.
-
     Args:
         num_qubits: The number of qubits for the MPS.
     """

src/qibotn/mps_utils.py

@@ -2,11 +2,10 @@ import cupy as cp
 from cuquantum import contract
 from cuquantum.cutensornet.experimental import contract_decompose
 
-# Reference: https://github.com/NVIDIA/cuQuantum/blob/main/python/samples/cutensornet/tn_algorithms/mps_algorithms.ipynb
 
 
 def initial(num_qubits, dtype):
-    r"""Generate the MPS with an initial state of :math:`\ket{00...00}`"""
+    """Generate the MPS with an initial state of :math:`\ket{00...00}`"""
     state_tensor = cp.asarray([1, 0], dtype=dtype).reshape(1, 2, 1)
     mps_tensors = [state_tensor] * num_qubits
     return mps_tensors
@@ -28,6 +27,8 @@ def mps_site_right_swap(mps_tensors, i, **kwargs):
 def apply_gate(mps_tensors, gate, qubits, **kwargs):
     """Apply the gate operand to the MPS tensors in-place.
     
+    # Reference: https://github.com/NVIDIA/cuQuantum/blob/main/python/samples/cutensornet/tn_algorithms/mps_algorithms.ipynb
+    
     Args:
         mps_tensors: A list of rank-3 ndarray-like tensor objects.
             The indices of the ith tensor are expected to be the bonding index to the i-1 tensor,