tensorcircuit-nightly 1.3.0.dev20250728__py3-none-any.whl → 1.4.0.dev20251103__py3-none-any.whl

tensorcircuit/fgs.py

@@ -28,7 +28,6 @@ def onehot_matrix(i: int, j: int, N: int) -> Tensor:
 
 # TODO(@refraction-ray): efficiency benchmark with jit
 # TODO(@refraction-ray): FGS mixed state support?
-# TODO(@refraction-ray): overlap?
 # TODO(@refraction-ray): fermionic logarithmic negativity
 
 
@@ -227,7 +226,7 @@ class FGSSimulator:
         return self.alpha
 
     def get_cmatrix(self, now_i: bool = True, now_j: bool = True) -> Tensor:
-        """
+        r"""
         Calculates the correlation matrix.
 
         The correlation matrix is defined as :math:`C_{ij} = \langle c_i^\dagger c_j \rangle`.
@@ -509,7 +508,7 @@ class FGSSimulator:
 
     @staticmethod
     def hopping(chi: Tensor, i: int, j: int, L: int) -> Tensor:
-        """
+        r"""
         Constructs the hopping Hamiltonian between two sites.
 
         The hopping Hamiltonian is given by :math:`\chi c_i^\dagger c_j + h.c.`.
@@ -550,7 +549,7 @@ class FGSSimulator:
 
     @staticmethod
     def chemical_potential(chi: Tensor, i: int, L: int) -> Tensor:
-        """
+        r"""
         Constructs the chemical potential Hamiltonian for a single site.
 
         The chemical potential Hamiltonian is given by :math:`\chi c_i^\dagger c_i`.
@@ -572,7 +571,7 @@ class FGSSimulator:
 
     @staticmethod
     def sc_pairing(chi: Tensor, i: int, j: int, L: int) -> Tensor:
-        """
+        r"""
         Constructs the superconducting pairing Hamiltonian between two sites.
 
         The superconducting pairing Hamiltonian is given by :math:`\chi c_i^\dagger c_j^\dagger + h.c.`.
@@ -637,7 +636,7 @@ class FGSSimulator:
         self.evol_ihamiltonian(self.chemical_potential(chi, i, self.L))
 
     def get_bogoliubov_uv(self) -> Tuple[Tensor, Tensor]:
-        """
+        r"""
         Returns the u and v matrices of the Bogoliubov transformation.
 
         The Bogoliubov transformation is defined as:

tensorcircuit/gates.py

@@ -34,6 +34,12 @@ one_state = np.array([0.0, 1.0], dtype=npdtype)
 plus_state = 1.0 / np.sqrt(2) * (zero_state + one_state)
 minus_state = 1.0 / np.sqrt(2) * (zero_state - one_state)
 
+# Common elements as np.ndarray objects
+_i00 = np.array([[1.0, 0.0], [0.0, 0.0]])
+_i01 = np.array([[0.0, 1.0], [0.0, 0.0]])
+_i10 = np.array([[0.0, 0.0], [1.0, 0.0]])
+_i11 = np.array([[0.0, 0.0], [0.0, 1.0]])
+
 # Common single qubit gates as np.ndarray objects
 _h_matrix = 1 / np.sqrt(2) * np.array([[1.0, 1.0], [1.0, -1.0]])
 _i_matrix = np.array([[1.0, 0.0], [0.0, 1.0]])
@@ -229,7 +235,7 @@ def num_to_tensor(*num: Union[float, Tensor], dtype: Optional[str] = None) -> An
     # TODO(@YHPeter): fix __doc__ for same function with different names
 
     l = []
-    if not dtype:
+    if dtype is None:
         dtype = dtypestr
     for n in num:
         if not backend.is_tensor(n):
@@ -245,7 +251,7 @@ array_to_tensor = num_to_tensor
 
 
 def gate_wrapper(m: Tensor, n: Optional[str] = None) -> Gate:
-    if not n:
+    if n is None:
         n = "unknowngate"
     m = m.astype(npdtype)
     return Gate(deepcopy(m), name=n)
@@ -255,7 +261,7 @@ class GateF:
     def __init__(
         self, m: Tensor, n: Optional[str] = None, ctrl: Optional[List[int]] = None
     ):
-        if not n:
+        if n is None:
             n = "unknowngate"
         self.m = m
         self.n = n
@@ -310,7 +316,7 @@ class GateF:
 
             return Gate(cu, name="c" + self.n)
 
-        if not self.ctrl:
+        if self.ctrl is None:
             ctrl = [1]
         else:
             ctrl = [1] + self.ctrl
@@ -330,7 +336,7 @@ class GateF:
             # TODO(@refraction-ray): ctrl convention to be finally determined
             return Gate(ocu, name="o" + self.n)
 
-        if not self.ctrl:
+        if self.ctrl is None:
             ctrl = [0]
         else:
             ctrl = [0] + self.ctrl
@@ -349,7 +355,7 @@ class GateVF(GateF):
         n: Optional[str] = None,
         ctrl: Optional[List[int]] = None,
     ):
-        if not n:
+        if n is None:
             n = "unknowngate"
         self.f = f
         self.n = n
@@ -483,7 +489,7 @@ def phase_gate(theta: float = 0) -> Gate:
     :rtype: Gate
     """
     theta = array_to_tensor(theta)
-    i00, i11 = array_to_tensor(np.array([[1, 0], [0, 0]]), np.array([[0, 0], [0, 1]]))
+    i00, i11 = array_to_tensor(_i00, _i11)
     unitary = i00 + backend.exp(1.0j * theta) * i11
     return Gate(unitary)
 
@@ -512,7 +518,7 @@ def get_u_parameter(m: Tensor) -> Tuple[float, float, float]:
     return theta, phi, lbd
 
 
-def u_gate(theta: float = 0, phi: float = 0, lbd: float = 0) -> Gate:
+def u_gate(theta: float = 0.0, phi: float = 0.0, lbd: float = 0.0) -> Gate:
     r"""
     IBMQ U gate following the converntion of OpenQASM3.0.
     See `OpenQASM doc <https://openqasm.com/language/gates.html#built-in-gates>`_
@@ -533,12 +539,7 @@ def u_gate(theta: float = 0, phi: float = 0, lbd: float = 0) -> Gate:
     :rtype: Gate
     """
     theta, phi, lbd = array_to_tensor(theta, phi, lbd)
-    i00, i01, i10, i11 = array_to_tensor(
-        np.array([[1, 0], [0, 0]]),
-        np.array([[0, 1], [0, 0]]),
-        np.array([[0, 0], [1, 0]]),
-        np.array([[0, 0], [0, 1]]),
-    )
+    i00, i01, i10, i11 = array_to_tensor(_i00, _i01, _i10, _i11)
     unitary = (
         backend.cos(theta / 2) * i00
         - backend.exp(1.0j * lbd) * backend.sin(theta / 2) * i01
@@ -548,7 +549,7 @@ def u_gate(theta: float = 0, phi: float = 0, lbd: float = 0) -> Gate:
     return Gate(unitary)
 
 
-def r_gate(theta: float = 0, alpha: float = 0, phi: float = 0) -> Gate:
+def r_gate(theta: float = 0.0, alpha: float = 0.0, phi: float = 0.0) -> Gate:
     r"""
     General single qubit rotation gate
 
@@ -582,7 +583,7 @@ def r_gate(theta: float = 0, alpha: float = 0, phi: float = 0) -> Gate:
 # r = r_gate
 
 
-def rx_gate(theta: float = 0) -> Gate:
+def rx_gate(theta: float = 0.0) -> Gate:
     r"""
     Rotation gate along :math:`x` axis.
 
@@ -603,7 +604,7 @@ def rx_gate(theta: float = 0) -> Gate:
 # rx = rx_gate
 
 
-def ry_gate(theta: float = 0) -> Gate:
+def ry_gate(theta: float = 0.0) -> Gate:
     r"""
     Rotation gate along :math:`y` axis.
 
@@ -624,7 +625,7 @@ def ry_gate(theta: float = 0) -> Gate:
 # ry = ry_gate
 
 
-def rz_gate(theta: float = 0) -> Gate:
+def rz_gate(theta: float = 0.0) -> Gate:
     r"""
     Rotation gate along :math:`z` axis.
 
@@ -645,7 +646,7 @@ def rz_gate(theta: float = 0) -> Gate:
 # rz = rz_gate
 
 
-def rgate_theoretical(theta: float = 0, alpha: float = 0, phi: float = 0) -> Gate:
+def rgate_theoretical(theta: float = 0.0, alpha: float = 0.0, phi: float = 0.0) -> Gate:
     r"""
     Rotation gate implemented by matrix exponential. The output is the same as `rgate`.
 
@@ -723,7 +724,7 @@ def iswap_gate(theta: float = 1.0) -> Gate:
 # iswap = iswap_gate
 
 
-def cr_gate(theta: float = 0, alpha: float = 0, phi: float = 0) -> Gate:
+def cr_gate(theta: float = 0.0, alpha: float = 0.0, phi: float = 0.0) -> Gate:
     r"""
     Controlled rotation gate. When the control qubit is 1, `rgate` is applied to the target qubit.
 
@@ -775,7 +776,7 @@ def random_two_qubit_gate() -> Gate:
     return Gate(deepcopy(unitary), name="R2Q")
 
 
-def any_gate(unitary: Tensor, name: str = "any") -> Gate:
+def any_gate(unitary: Tensor, name: str = "any", dim: Optional[int] = None) -> Gate:
     """
     Note one should provide the gate with properly reshaped.
 
@@ -783,6 +784,8 @@ def any_gate(unitary: Tensor, name: str = "any") -> Gate:
     :type unitary: Tensor
     :param name: The name of the gate.
     :type name: str
+    :param dim: The dimension of the gate.
+    :type dim: int
     :return: the resulted gate
     :rtype: Gate
     """
@@ -791,7 +794,10 @@ def any_gate(unitary: Tensor, name: str = "any") -> Gate:
         unitary.tensor = backend.cast(unitary.tensor, dtypestr)
         return unitary
     unitary = backend.cast(unitary, dtypestr)
-    unitary = backend.reshape2(unitary)
+    if dim is None or dim == 2:
+        unitary = backend.reshape2(unitary)
+    else:
+        unitary = backend.reshaped(unitary, dim)
     # nleg = int(np.log2(backend.sizen(unitary)))
     # unitary = backend.reshape(unitary, [2 for _ in range(nleg)])
     return Gate(unitary, name=name)
@@ -864,6 +870,43 @@ def exponential_gate_unity(
     return Gate(mat, name="exp1-" + name)
 
 
+def su4_gate(theta: Tensor, name: str = "su(4)") -> Gate:
+    r"""
+    Two-qubit general SU(4) gate.
+
+    :param theta: the angle tensor (15 components) of the gate.
+    :type theta: Tensor
+    :param name: the name of the gate.
+    :type name: str
+    :return: a gate object.
+    :rtype: Gate
+    """
+    theta = num_to_tensor(theta)
+    pauli_ops = array_to_tensor(
+        _ix_matrix,
+        _iy_matrix,
+        _iz_matrix,
+        _xi_matrix,
+        _xx_matrix,
+        _xy_matrix,
+        _xz_matrix,
+        _yi_matrix,
+        _yx_matrix,
+        _yy_matrix,
+        _yz_matrix,
+        _zi_matrix,
+        _zx_matrix,
+        _zy_matrix,
+        _zz_matrix,
+    )
+    generator = backend.sum(
+        backend.stack([theta[i] * pauli_ops[i] for i in range(15)]), axis=0
+    )
+    mat = backend.expm(-1j * generator)
+    mat = backend.reshape2(mat)
+    return Gate(mat, name=name)
+
+
 exp1_gate = exponential_gate_unity
 # exp1 = exponential_gate_unity
 rzz_gate = partial(exp1_gate, unitary=_zz_matrix, half=True)
@@ -968,6 +1011,7 @@ def meta_vgate() -> None:
         "rzz",
         "rxx",
         "ryy",
+        "su4",
     ]:
         for funcname in [f, f + "gate"]:
             setattr(thismodule, funcname, GateVF(getattr(thismodule, f + "_gate"), f))

tensorcircuit/keras.py

@@ -24,7 +24,7 @@ class QuantumLayer(Layer):  # type: ignore
         initializer: Union[Text, Sequence[Text]] = "glorot_uniform",
         constraint: Optional[Union[Text, Sequence[Text]]] = None,
         regularizer: Optional[Union[Text, Sequence[Text]]] = None,
-        **kwargs: Any
+        **kwargs: Any,
     ) -> None:
         """
         `QuantumLayer` wraps the quantum function `f` as a `keras.Layer`
@@ -103,7 +103,7 @@ class QuantumLayer(Layer):  # type: ignore
         inputs: tf.Tensor,
         training: Optional[bool] = None,
         mask: Optional[tf.Tensor] = None,
-        **kwargs: Any
+        **kwargs: Any,
     ) -> tf.Tensor:
         # input_shape = list(inputs.shape)
         # inputs = tf.reshape(inputs, (-1, input_shape[-1]))
@@ -154,7 +154,7 @@ class HardwareLayer(QuantumLayer):
         inputs: tf.Tensor,
         training: Optional[bool] = None,
         mask: Optional[tf.Tensor] = None,
-        **kwargs: Any
+        **kwargs: Any,
     ) -> tf.Tensor:
         if inputs is None:  # not possible
             result = self.f(*self.pqc_weights, **kwargs)

tensorcircuit/mpscircuit.py

@@ -4,21 +4,25 @@ Quantum circuit: MPS state simulator
 
 # pylint: disable=invalid-name
 
-from functools import reduce
+from functools import reduce, partial
 from typing import Any, List, Optional, Sequence, Tuple, Dict, Union
 from copy import copy
+import logging
+import types
 
 import numpy as np
 import tensornetwork as tn
-from tensorcircuit.quantum import QuOperator, QuVector
 
 from . import gates
 from .cons import backend, npdtype, contractor, rdtypestr, dtypestr
+from .quantum import QuOperator, QuVector, extract_tensors_from_qop, _decode_basis_label
 from .mps_base import FiniteMPS
 from .abstractcircuit import AbstractCircuit
+from .utils import arg_alias
 
 Gate = gates.Gate
 Tensor = Any
+logger = logging.getLogger(__name__)
 
 
 def split_tensor(
@@ -77,6 +81,10 @@ class MPSCircuit(AbstractCircuit):
 
     is_mps = True
 
+    @partial(
+        arg_alias,
+        alias_dict={"wavefunction": ["inputs"]},
+    )
     def __init__(
         self,
         nqubits: int,
@@ -84,12 +92,16 @@ class MPSCircuit(AbstractCircuit):
         tensors: Optional[Sequence[Tensor]] = None,
         wavefunction: Optional[Union[QuVector, Tensor]] = None,
         split: Optional[Dict[str, Any]] = None,
+        dim: Optional[int] = None,
     ) -> None:
         """
         MPSCircuit object based on state simulator.
+        Do not use this class with d!=2 directly
 
         :param nqubits: The number of qubits in the circuit.
         :type nqubits: int
+        :param dim: The local Hilbert space dimension per site. Qudit is supported for 2 <= d <= 36.
+        :type dim: If None, the dimension of the circuit will be `2`, which is a qubit system.
         :param center_position: The center position of MPS, default to 0
         :type center_position: int, optional
         :param tensors: If not None, the initial state of the circuit is taken as ``tensors``
@@ -102,6 +114,7 @@ class MPSCircuit(AbstractCircuit):
         :param split: Split rules
         :type split: Any
         """
+        self._d = 2 if dim is None else dim
         self.circuit_param = {
             "nqubits": nqubits,
             "center_position": center_position,
@@ -118,8 +131,21 @@ class MPSCircuit(AbstractCircuit):
             ), "tensors and wavefunction cannot be used at input simutaneously"
             # TODO(@SUSYUSTC): find better way to address QuVector
             if isinstance(wavefunction, QuVector):
-                wavefunction = wavefunction.eval()
-            tensors = self.wavefunction_to_tensors(wavefunction, split=self.split)
+                try:
+                    nodes, is_mps, _ = extract_tensors_from_qop(wavefunction)
+                    if not is_mps:
+                        raise ValueError("wavefunction is not a valid MPS")
+                    tensors = [node.tensor for node in nodes]
+                except ValueError as e:
+                    logger.warning(repr(e))
+                    wavefunction = wavefunction.eval()
+                    tensors = self.wavefunction_to_tensors(
+                        wavefunction, split=self.split
+                    )
+            else:  # full wavefunction
+                tensors = self.wavefunction_to_tensors(
+                    wavefunction, dim_phys=self._d, split=self.split
+                )
             assert len(tensors) == nqubits
             self._mps = FiniteMPS(tensors, canonicalize=False)
             self._mps.center_position = 0
@@ -133,8 +159,13 @@ class MPSCircuit(AbstractCircuit):
                 self._mps = FiniteMPS(tensors, canonicalize=True, center_position=0)
         else:
             tensors = [
-                np.array([1.0, 0.0], dtype=npdtype)[None, :, None]
-                for i in range(nqubits)
+                np.concatenate(
+                    [
+                        np.array([1.0], dtype=npdtype),
+                        np.zeros((self._d - 1,), dtype=npdtype),
+                    ]
+                )[None, :, None]
+                for _ in range(nqubits)
             ]
             self._mps = FiniteMPS(tensors, canonicalize=False)
             if center_position is not None:
@@ -352,42 +383,51 @@ class MPSCircuit(AbstractCircuit):
         #      b
 
         # index must be ordered
-        assert np.all(np.diff(index) > 0)
-        index_left = np.min(index)
+        if len(index) == 0:
+            raise ValueError("`index` must contain at least one site.")
+        if not all(index[i] < index[i + 1] for i in range(len(index) - 1)):
+            raise ValueError("`index` must be strictly increasing.")
+
+        index_left = int(np.min(index))
         if isinstance(gate, tn.Node):
             gate = backend.copy(gate.tensor)
-        index = np.array(index) - index_left
+
         nindex = len(index)
+        in_dims = tuple(backend.shape_tuple(gate))[:nindex]
+        dim = int(in_dims[0])
+        dim_phys_mpo = dim * dim
+        gate = backend.reshape(gate, (dim,) * nindex + (dim,) * nindex)
         # transform gate from (in1, in2, ..., out1, out2 ...) to
         # (in1, out1, in2, out2, ...)
-        order = tuple(np.arange(2 * nindex).reshape((2, nindex)).T.flatten())
-        shape = (4,) * nindex
-        gate = backend.reshape(backend.transpose(gate, order), shape)
-        argsort = np.argsort(index)
+        order = tuple(np.arange(2 * nindex).reshape(2, nindex).T.flatten().tolist())
+        gate = backend.transpose(gate, order)
         # reorder the gate according to the site positions
-        gate = backend.transpose(gate, tuple(argsort))
-        index = index[argsort]  # type: ignore
+        gate = backend.reshape(gate, (dim_phys_mpo,) * nindex)
         # split the gate into tensors assuming they are adjacent
-        main_tensors = cls.wavefunction_to_tensors(gate, dim_phys=4, norm=False)
+        main_tensors = cls.wavefunction_to_tensors(
+            gate, dim_phys=dim_phys_mpo, norm=False
+        )
         # each tensor is in shape of (i, a, b, j)
-        tensors = []
-        previous_i = None
-        for i, main_tensor in zip(index, main_tensors):
-            # insert identites in the middle
+        tensors: list[Tensor] = []
+        previous_i: Optional[int] = None
+        index_arr = np.array(index, dtype=int) - index_left
+
+        for i, main_tensor in zip(index_arr, main_tensors):
             if previous_i is not None:
-                for _ in range(previous_i + 1, i):
-                    bond_dim = tensors[-1].shape[-1]
-                    I = (
-                        np.eye(bond_dim * 2)
-                        .reshape((bond_dim, 2, bond_dim, 2))
-                        .transpose((0, 1, 3, 2))
-                        .astype(dtypestr)
+                for _gap_site in range(int(previous_i) + 1, int(i)):
+                    bond_dim = int(backend.shape_tuple(tensors[-1])[-1])
+                    eye2d = backend.eye(
+                        bond_dim * dim, dtype=backend.dtype(tensors[-1])
                     )
-                    tensors.append(backend.convert_to_tensor(I))
-            nleft, _, nright = main_tensor.shape
-            tensor = backend.reshape(main_tensor, (nleft, 2, 2, nright))
+                    I4 = backend.reshape(eye2d, (bond_dim, dim, bond_dim, dim))
+                    I4 = backend.transpose(I4, (0, 1, 3, 2))
+                    tensors.append(I4)
+
+            nleft, _, nright = backend.shape_tuple(main_tensor)
+            tensor = backend.reshape(main_tensor, (int(nleft), dim, dim, int(nright)))
             tensors.append(tensor)
-            previous_i = i
+            previous_i = int(i)
+
         return tensors, index_left
 
     @classmethod
@@ -430,15 +470,15 @@ class MPSCircuit(AbstractCircuit):
         """
         if split is None:
             split = {}
-        ni = tensor_left.shape[0]
-        nk = tensor_right.shape[-1]
+        ni, di = tensor_left.shape[0], tensor_right.shape[1]
+        nk, dk = tensor_right.shape[-1], tensor_right.shape[-2]
         T = backend.einsum("iaj,jbk->iabk", tensor_left, tensor_right)
-        T = backend.reshape(T, (ni * 2, nk * 2))
+        T = backend.reshape(T, (ni * di, nk * dk))
         new_tensor_left, new_tensor_right = split_tensor(
             T, center_left=center_left, split=split
         )
-        new_tensor_left = backend.reshape(new_tensor_left, (ni, 2, -1))
-        new_tensor_right = backend.reshape(new_tensor_right, (-1, 2, nk))
+        new_tensor_left = backend.reshape(new_tensor_left, (ni, di, -1))
+        new_tensor_right = backend.reshape(new_tensor_right, (-1, dk, nk))
         return new_tensor_left, new_tensor_right
 
     def reduce_dimension(
@@ -532,10 +572,11 @@ class MPSCircuit(AbstractCircuit):
         for i, idx in zip(i_list, idx_list):
             O = tensors[i]
             T = self._mps.tensors[idx]
-            ni, _, _, nj = O.shape
+            ni, d_in, _, nj = O.shape
             nk, _, nl = T.shape
             OT = backend.einsum("iabj,kbl->ikajl", O, T)
-            OT = backend.reshape(OT, (ni * nk, 2, nj * nl))
+            OT = backend.reshape(OT, (ni * nk, d_in, nj * nl))
+
             self._mps.tensors[idx] = OT
 
         # canonicalize
@@ -642,8 +683,7 @@ class MPSCircuit(AbstractCircuit):
         :type keep: int, optional
         """
         # normalization not guaranteed
-        assert keep in [0, 1]
-        gate = backend.zeros((2, 2), dtype=dtypestr)
+        gate = backend.zeros((self._d, self._d), dtype=dtypestr)
         gate = backend.scatter(
             gate,
             backend.convert_to_tensor([[keep, keep]]),
@@ -674,7 +714,7 @@ class MPSCircuit(AbstractCircuit):
     def wavefunction_to_tensors(
         cls,
         wavefunction: Tensor,
-        dim_phys: int = 2,
+        dim_phys: Optional[int] = None,
         norm: bool = True,
         split: Optional[Dict[str, Any]] = None,
     ) -> List[Tensor]:
@@ -692,6 +732,7 @@ class MPSCircuit(AbstractCircuit):
         :return: The tensors
         :rtype: List[Tensor]
         """
+        dim_phys = dim_phys if dim_phys is not None else 2
         if split is None:
             split = {}
         wavefunction = backend.reshape(wavefunction, (-1, 1))
@@ -750,10 +791,16 @@ class MPSCircuit(AbstractCircuit):
         for key in vars(self):
             if key == "_mps":
                 continue
-            if backend.is_tensor(info[key]):
-                copied_value = backend.copy(info[key])
+            val = info[key]
+            if backend.is_tensor(val):
+                copied_value = backend.copy(val)
+            elif isinstance(val, types.ModuleType):
+                copied_value = val
             else:
-                copied_value = copy(info[key])
+                try:
+                    copied_value = copy(val)
+                except TypeError:
+                    copied_value = val
             setattr(result, key, copied_value)
         return result
 
@@ -797,7 +844,8 @@ class MPSCircuit(AbstractCircuit):
 
     def amplitude(self, l: str) -> Tensor:
         assert len(l) == self._nqubits
-        tensors = [self._mps.tensors[i][:, int(s), :] for i, s in enumerate(l)]
+        idx_list = _decode_basis_label(l, n=self._nqubits, dim=self._d)
+        tensors = [self._mps.tensors[i][:, idx, :] for i, idx in enumerate(idx_list)]
         return reduce(backend.matmul, tensors)[0, 0]
 
     def proj_with_mps(self, other: "MPSCircuit", conj: bool = True) -> Tensor:
@@ -855,6 +903,7 @@ class MPSCircuit(AbstractCircuit):
 
         mps = self.__class__(
             nqubits,
+            dim=self._d,
             tensors=tensors,
             center_position=center_position,
             split=self.split.copy(),
@@ -982,36 +1031,35 @@ class MPSCircuit(AbstractCircuit):
         # set the center to the left side, then gradually move to the right and do measurement at sites
         """
         mps = self.copy()
-        up = backend.convert_to_tensor(np.array([1, 0]).astype(dtypestr))
-        down = backend.convert_to_tensor(np.array([0, 1]).astype(dtypestr))
 
         p = 1.0
         p = backend.convert_to_tensor(p)
         p = backend.cast(p, dtype=rdtypestr)
-        sample = []
+        sample: Tensor = []
         for k, site in enumerate(index):
             mps.position(site)
-            # do measurement
             tensor = mps._mps.tensors[site]
             ps = backend.real(
                 backend.einsum("iaj,iaj->a", tensor, backend.conj(tensor))
             )
             ps /= backend.sum(ps)
-            pu = ps[0]
             if status is None:
-                r = backend.implicit_randu()[0]
+                outcome = backend.implicit_randc(
+                    self._d, shape=1, p=backend.cast(ps, rdtypestr)
+                )[0]
             else:
-                r = status[k]
-            r = backend.real(backend.cast(r, dtypestr))
-            eps = 0.31415926 * 1e-12
-            sign = backend.sign(r - pu + eps) / 2 + 0.5  # in case status is exactly 0.5
-            sign = backend.convert_to_tensor(sign)
-            sign = backend.cast(sign, dtype=rdtypestr)
-            sign_complex = backend.cast(sign, dtypestr)
-            sample.append(sign_complex)
-            p = p * (pu * (-1) ** sign + sign)
-            m = (1 - sign_complex) * up + sign_complex * down
+                one_r = backend.cast(backend.convert_to_tensor(1.0), rdtypestr)
+                st = backend.cast(status[k : k + 1], rdtypestr)
+                ind = backend.probability_sample(
+                    shots=1, p=backend.cast(ps, rdtypestr), status=one_r - st
+                )
+                outcome = backend.cast(ind[0], "int32")
+
+            p = p * ps[outcome]
+            basis = backend.convert_to_tensor(np.eye(self._d).astype(dtypestr))
+            m = basis[outcome]
             mps._mps.tensors[site] = backend.einsum("iaj,a->ij", tensor, m)[:, None, :]
+            sample.append(outcome)
         sample = backend.stack(sample)
         sample = backend.real(sample)
         if with_prob: