From 935001af2560f40d9ff784539fcd9f04cf982711 Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Wed, 27 Aug 2025 10:12:24 +0800
Subject: [PATCH 01/55] Fixed a potential issue in the qudit system that could
 have caused latent errors.

---
 tensorcircuit/simplify.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tensorcircuit/simplify.py b/tensorcircuit/simplify.py
index 1cbcfc6a..917e6d54 100644
--- a/tensorcircuit/simplify.py
+++ b/tensorcircuit/simplify.py
@@ -121,7 +121,9 @@ def _split_two_qubit_gate(
     if fixed_choice == 2:  # swap one
         return n3, n4, True  # swap
     s2 = n3.tensor.shape[-1]
-    if (s1 >= 4) and (s2 >= 4):
+    if (s1 >= n[0].dimension * n[2].dimension) and (
+        s2 >= n[1].dimension * n[3].dimension
+    ):
         # jax jit unspport split_node with trun_err anyway
         # tf function doesn't work either, though I believe it may work on tf side
         # CANNOT DONE(@refraction-ray): tf.function version with trun_err set

From db7b20292f496c3ef1496180a643d83386daa081 Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Wed, 27 Aug 2025 10:14:09 +0800
Subject: [PATCH 02/55] Add variable _ALPHBET for representing the states in
 qudit systems.

---
 tensorcircuit/cons.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorcircuit/cons.py b/tensorcircuit/cons.py
index c7db79ac..660b1c2a 100644
--- a/tensorcircuit/cons.py
+++ b/tensorcircuit/cons.py
@@ -63,6 +63,7 @@ def sorted_edges(edges: Iterator[tn.Edge]) -> List[tn.Edge]:
 npdtype = np.complex64
 backend: NumpyBackend = get_backend("numpy")
 contractor = tn.contractors.auto
+_ALPHABET = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ"
 # these above lines are just for mypy, it is not very good at evaluating runtime object
 
 

From 81cca4f58855fa879440f9da22ead2cc87ff403e Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Wed, 27 Aug 2025 10:15:38 +0800
Subject: [PATCH 03/55] Optimized functions in counts, make them avaliable for
 qudit systems. While a qudit in |10> state, it will be shown as `A` from
 _ALPHBET.

---
 tensorcircuit/results/counts.py | 113 +++++++++++++++++++++++---------
 1 file changed, 83 insertions(+), 30 deletions(-)

diff --git a/tensorcircuit/results/counts.py b/tensorcircuit/results/counts.py
index 5bf0e870..b38dada8 100644
--- a/tensorcircuit/results/counts.py
+++ b/tensorcircuit/results/counts.py
@@ -2,10 +2,11 @@
 dict related functionalities
 """
 
-from typing import Any, Dict, Optional, Sequence
+from typing import Any, Dict, Optional, Sequence, List
 
 import numpy as np
 
+from ..cons import _ALPHABET
 
 Tensor = Any
 ct = Dict[str, int]
@@ -91,7 +92,13 @@ def marginal_count(count: ct, keep_list: Sequence[int]) -> ct:
 
 def count2vec(count: ct, normalization: bool = True) -> Tensor:
     """
-    Convert count dictionary to probability vector.
+    Convert a dictionary of counts (with string keys) to a probability/count vector.
+
+    Support:
+      - base-d string (d <= 36), characters taken from 0-9A-Z (case-insensitive)
+        For example:
+          qubit: '0101'
+          qudit: '012' or '09A' (A represents 10, which means [0, 9, 10])
 
     :param count: A dictionary mapping bit strings to counts
     :type count: ct
@@ -105,44 +112,90 @@ def count2vec(count: ct, normalization: bool = True) -> Tensor:
     >>> count2vec({"00": 2, "10": 3, "11": 5})
     array([0.2, 0. , 0.3, 0.5])
     """
-    nqubit = len(list(count.keys())[0])
-    probability = [0] * 2**nqubit
-    shots = sum([v for k, v in count.items()])
+    if not count:
+        return np.array([], dtype=float)
+
+    sample_key = next(iter(count)).upper()
+    n = len(sample_key)
+    d = 0
+    for k in count:
+        s = k.upper()
+        if len(s) != n:
+            raise ValueError(
+                f"The length of all keys should be the same ({n}), received '{k}'."
+            )
+        for ch in s:
+            if ch not in _ALPHABET:
+                raise ValueError(
+                    f"Key '{k}' contains illegal character '{ch}' (only 0-9A-Z are allowed)."
+                )
+            d = max(d, _ALPHABET.index(ch) + 1)
+    if d < 2:
+        raise ValueError(f"Inferred local dimension d={d} is illegal (must be >=2).")
+
+    def parse_key(_k: str) -> List[int]:
+        return [_ALPHABET.index(_ch) for _ch in _k.upper()]
+
+    size = d**n
+    prob = np.zeros(size, dtype=float)
+    shots = float(sum(count.values())) if normalization else 1.0
+    if shots == 0:
+        return prob
+
+    powers = [d**p for p in range(n)][::-1]
     for k, v in count.items():
-        if normalization is True:
-            v /= shots  # type: ignore
-        probability[int(k, 2)] = v
-    return np.array(probability)
+        digits = parse_key(k)
+        idx = sum(dig * p for dig, p in zip(digits, powers))
+        prob[idx] = (v / shots) if normalization else v
+
+    return prob
 
 
 def vec2count(vec: Tensor, prune: bool = False) -> ct:
     """
-    Convert probability vector to count dictionary.
-
-    :param vec: Probability vector
-    :type vec: Tensor
-    :param prune: Whether to remove near-zero probabilities, defaults to False
-    :type prune: bool, optional
-    :return: Count dictionary
-    :rtype: ct
-
-    :Example:
+    Map a count/probability vector of length D to a dictionary with base-d string keys (0-9A-Z).
+    Only generate string keys when d ≤ 36; if d is inferred to be > 36, raise a NotImplementedError.
 
-    >>> vec2count(np.array([0.2, 0.3, 0.1, 0.4]))
-    {'00': 0.2, '01': 0.3, '10': 0.1, '11': 0.4}
+    :param vec: A one-dimensional vector of length D = d**n
+    :param prune: Whether to prune near-zero elements (threshold 1e-8)
+    :return: {base-d string key: value}, key length n
     """
-    from ..quantum import count_vector2dict
+    from ..quantum import count_vector2dict, _infer_num_sites
 
     if isinstance(vec, list):
         vec = np.array(vec)
-    n = int(np.log(vec.shape[0]) / np.log(2) + 1e-9)
-    c = count_vector2dict(vec, n, key="bin")
-    if prune is True:
-        nc = c.copy()
-        for k, v in c.items():
-            if np.abs(v) < 1e-8:
-                del nc[k]
-        return nc
+    vec = np.asarray(vec)
+    if vec.ndim != 1:
+        raise ValueError("vec2count expects a one-dimensional vector.")
+
+    D = int(vec.shape[0])
+    if D <= 0:
+        return {}
+
+    def _is_power_of_two(x: int) -> bool:
+        return x > 0 and (x & (x - 1)) == 0
+
+    if _is_power_of_two(D):
+        n = int(np.log(D) / np.log(2) + 1e-9)
+        d: Optional[int] = 2
+    else:
+        d = n = None
+        upper = int(np.sqrt(D)) + 1
+        for d_try in range(2, max(upper, 3)):
+            try:
+                n_try = _infer_num_sites(D, d_try)
+            except ValueError:
+                continue
+            d, n = d_try, n_try
+            break
+        if d is None:
+            d, n = D, 1
+
+    c: ct = count_vector2dict(vec, n, key="bin", d=d)  # type: ignore
+
+    if prune:
+        c = {k: v for k, v in c.items() if np.abs(v) >= 1e-8}
+
     return c
 
 

From 52fbf669860fd6f5cdf7027e20123061d786bc05 Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Wed, 27 Aug 2025 10:19:30 +0800
Subject: [PATCH 04/55] Removed hard codes in quantum.py (corresponding to
 qubit system.), all the functions are now avaliable for any d-dimensional
 systems. Add _infer_num_sites() for inferring the site number in
 d-dimensional system.

---
 tensorcircuit/quantum.py | 430 +++++++++++++++++++++++++--------------
 1 file changed, 280 insertions(+), 150 deletions(-)

diff --git a/tensorcircuit/quantum.py b/tensorcircuit/quantum.py
index d28e550c..a6b60a66 100644
--- a/tensorcircuit/quantum.py
+++ b/tensorcircuit/quantum.py
@@ -56,6 +56,30 @@ def get_all_nodes(edges: Iterable[Edge]) -> List[Node]:
     return nodes
 
 
+def _infer_num_sites(D: int, d: int) -> int:
+    """
+    Infer the number of sites (n) from a Hilbert space dimension D
+    and local dimension d, assuming D = d**n.
+
+    :param D: total Hilbert space dimension (int)
+    :param d: local dimension per site (int)
+    :return: n such that D == d**n
+    :raises ValueError: if D is not an exact power of d
+    """
+    if not (isinstance(D, int) and D > 0):
+        raise ValueError(f"D must be a positive integer, got {D}")
+    if not (isinstance(d, int) and d >= 2):
+        raise ValueError(f"d must be an integer >= 2, got {d}")
+
+    tmp, n = D, 0
+    while tmp % d == 0 and tmp > 1:
+        tmp //= d
+        n += 1
+    if tmp != 1:
+        raise ValueError(f"Dimension {D} is not a power of local dim {d}")
+    return n
+
+
 def _reachable(nodes: List[AbstractNode]) -> List[AbstractNode]:
     if not nodes:
         raise ValueError("Reachable requires at least 1 node.")
@@ -2150,7 +2174,10 @@ def entanglement_entropy(state: Tensor, cut: Union[int, List[int]]) -> Tensor:
 
 
 def reduced_wavefunction(
-    state: Tensor, cut: List[int], measure: Optional[List[int]] = None
+    state: Tensor,
+    cut: List[int],
+    measure: Optional[List[int]] = None,
+    d: Optional[int] = None,
 ) -> Tensor:
     """
     Compute the reduced wavefunction from the quantum state ``state``.
@@ -2165,20 +2192,22 @@ def reduced_wavefunction(
     :type measure: List[int]
     :return: _description_
     :rtype: Tensor
+    :param d: dimension of qudit system
+    :type d: int
     """
+    d = 2 if d is None else d
     if measure is None:
         measure = [0 for _ in cut]
-    s = backend.reshape2(state)
+    s = backend.reshaped(state, d)
     n = len(backend.shape_tuple(s))
     s_node = Gate(s)
     end_nodes = []
     for c, m in zip(cut, measure):
-        rt = backend.cast(backend.convert_to_tensor(1 - m), dtypestr) * backend.cast(
-            backend.convert_to_tensor(np.array([1.0, 0.0])), dtypestr
-        ) + backend.cast(backend.convert_to_tensor(m), dtypestr) * backend.cast(
-            backend.convert_to_tensor(np.array([0.0, 1.0])), dtypestr
+        oh = backend.cast(
+            backend.one_hot(backend.cast(backend.convert_to_tensor(m), "int32"), d),
+            dtypestr,
         )
-        end_node = Gate(rt)
+        end_node = Gate(backend.convert_to_tensor(oh))
         end_nodes.append(end_node)
         s_node[c] ^ end_node[0]
     new_node = contractor(
@@ -2193,8 +2222,9 @@ def reduced_density_matrix(
     cut: Union[int, List[int]],
     p: Optional[Tensor] = None,
     normalize: bool = True,
+    d: Optional[int] = None,
 ) -> Union[Tensor, QuOperator]:
-    """
+    r"""
     Compute the reduced density matrix from the quantum state ``state``.
 
     :param state: The quantum state in form of Tensor or QuOperator.
@@ -2206,8 +2236,12 @@ def reduced_density_matrix(
     :type p: Optional[Tensor]
     :return: The reduced density matrix.
     :rtype: Union[Tensor, QuOperator]
-    :normalize: if True, returns a trace 1 density matrix. Otherwise does not normalize.
+    :param normalize: if True, returns a trace 1 density matrix. Otherwise does not normalize.
+    :type normalize: bool
+    :param d: dimension of qudit system
+    :type d: int
     """
+    d = 2 if d is None else d
     if isinstance(cut, list) or isinstance(cut, tuple) or isinstance(cut, set):
         traceout = list(cut)
     else:
@@ -2220,21 +2254,19 @@ def reduced_density_matrix(
         return state.partial_trace(traceout)
     if len(state.shape) == 2 and state.shape[0] == state.shape[1]:
         # density operator
-        freedomexp = backend.sizen(state)
-        # traceout = sorted(traceout)[::-1]
-        freedom = int(np.log2(freedomexp) / 2)
-        # traceout2 = [i + freedom for i in traceout]
+        freedom = _infer_num_sites(state.shape[0], d)
         left = traceout + [i for i in range(freedom) if i not in traceout]
         right = [i + freedom for i in left]
-        rho = backend.reshape(state, [2 for _ in range(2 * freedom)])
+
+        rho = backend.reshape(state, [d] * (2 * freedom))
         rho = backend.transpose(rho, perm=left + right)
         rho = backend.reshape(
             rho,
             [
-                2 ** len(traceout),
-                2 ** (freedom - len(traceout)),
-                2 ** len(traceout),
-                2 ** (freedom - len(traceout)),
+                d ** len(traceout),
+                d ** (freedom - len(traceout)),
+                d ** len(traceout),
+                d ** (freedom - len(traceout)),
             ],
         )
         if p is None:
@@ -2247,20 +2279,20 @@ def reduced_density_matrix(
             p = backend.reshape(p, [-1])
             rho = backend.einsum("a,aiaj->ij", p, rho)
         rho = backend.reshape(
-            rho, [2 ** (freedom - len(traceout)), 2 ** (freedom - len(traceout))]
+            rho, [d ** (freedom - len(traceout)), d ** (freedom - len(traceout))]
         )
         if normalize:
             rho /= backend.trace(rho)
 
     else:
         w = state / backend.norm(state)
-        freedomexp = backend.sizen(state)
-        freedom = int(np.log(freedomexp) / np.log(2))
+        size = int(backend.sizen(state))
+        freedom = _infer_num_sites(size, d)
         perm = [i for i in range(freedom) if i not in traceout]
         perm = perm + traceout
-        w = backend.reshape(w, [2 for _ in range(freedom)])
+        w = backend.reshape(w, [d for _ in range(freedom)])
         w = backend.transpose(w, perm=perm)
-        w = backend.reshape(w, [-1, 2 ** len(traceout)])
+        w = backend.reshape(w, [-1, d ** len(traceout)])
         if p is None:
             rho = w @ backend.adjoint(w)
         else:
@@ -2403,7 +2435,9 @@ def truncated_free_energy(
 
 
 @op2tensor
-def partial_transpose(rho: Tensor, transposed_sites: List[int]) -> Tensor:
+def partial_transpose(
+    rho: Tensor, transposed_sites: List[int], d: Optional[int] = None
+) -> Tensor:
     """
     _summary_
 
@@ -2411,10 +2445,13 @@ def partial_transpose(rho: Tensor, transposed_sites: List[int]) -> Tensor:
     :type rho: Tensor
     :param transposed_sites: sites int list to be transposed
     :type transposed_sites: List[int]
+    :param d: dimension of qudit system
+    :type d: int
     :return: _description_
     :rtype: Tensor
     """
-    rho = backend.reshape2(rho)
+    d = 2 if d is None else d
+    rho = backend.reshaped(rho, d)
     rho_node = Gate(rho)
     n = len(rho.shape) // 2
     left_edges = []
@@ -2432,7 +2469,9 @@ def partial_transpose(rho: Tensor, transposed_sites: List[int]) -> Tensor:
 
 
 @op2tensor
-def entanglement_negativity(rho: Tensor, transposed_sites: List[int]) -> Tensor:
+def entanglement_negativity(
+    rho: Tensor, transposed_sites: List[int], d: Optional[int] = None
+) -> Tensor:
     """
     _summary_
 
@@ -2440,6 +2479,8 @@ def entanglement_negativity(rho: Tensor, transposed_sites: List[int]) -> Tensor:
     :type rho: Tensor
     :param transposed_sites: _description_
     :type transposed_sites: List[int]
+    :param d: dimension of qudit system
+    :type d: int
     :return: _description_
     :rtype: Tensor
     """
@@ -2450,7 +2491,9 @@ def entanglement_negativity(rho: Tensor, transposed_sites: List[int]) -> Tensor:
 
 
 @op2tensor
-def log_negativity(rho: Tensor, transposed_sites: List[int], base: str = "e") -> Tensor:
+def log_negativity(
+    rho: Tensor, transposed_sites: List[int], base: str = "e", d: Optional[int] = None
+) -> Tensor:
     """
     _summary_
 
@@ -2460,10 +2503,13 @@ def log_negativity(rho: Tensor, transposed_sites: List[int], base: str = "e") ->
     :type transposed_sites: List[int]
     :param base: whether use 2 based log or e based log, defaults to "e"
     :type base: str, optional
+    :param d: dimension of qudit system
+    :type d: int
     :return: _description_
     :rtype: Tensor
     """
-    rhot = partial_transpose(rho, transposed_sites)
+    d = 2 if d is None else d
+    rhot = partial_transpose(rho, transposed_sites, d)
     es = backend.eigvalsh(rhot)
     rhot_m = backend.sum(backend.abs(es))
     een = backend.log(rhot_m)
@@ -2549,7 +2595,9 @@ def double_state(h: Tensor, beta: float = 1) -> Tensor:
 
 
 @op2tensor
-def mutual_information(s: Tensor, cut: Union[int, List[int]]) -> Tensor:
+def mutual_information(
+    s: Tensor, cut: Union[int, List[int]], d: Optional[int] = None
+) -> Tensor:
     """
     Mutual information between AB subsystem described by ``cut``.
 
@@ -2557,9 +2605,12 @@ def mutual_information(s: Tensor, cut: Union[int, List[int]]) -> Tensor:
     :type s: Tensor
     :param cut: The AB subsystem.
     :type cut: Union[int, List[int]]
+    :param d: The diagonal matrix in form of Tensor.
+    :type d: Tensor
     :return: The mutual information between AB subsystem described by ``cut``.
     :rtype: Tensor
     """
+    d = 2 if d is None else d
     if isinstance(cut, list) or isinstance(cut, tuple) or isinstance(cut, set):
         traceout = list(cut)
     else:
@@ -2567,22 +2618,22 @@ def mutual_information(s: Tensor, cut: Union[int, List[int]]) -> Tensor:
 
     if len(s.shape) == 2 and s.shape[0] == s.shape[1]:
         # mixed state
-        n = int(np.log2(backend.sizen(s)) / 2)
+        n = _infer_num_sites(s.shape[0], d=d)
         hab = entropy(s)
 
         # subsystem a
-        rhoa = reduced_density_matrix(s, traceout)
+        rhoa = reduced_density_matrix(s, traceout, d=d)
         ha = entropy(rhoa)
 
         # need subsystem b as well
         other = tuple(i for i in range(n) if i not in traceout)
-        rhob = reduced_density_matrix(s, other)  # type: ignore
+        rhob = reduced_density_matrix(s, other, d=d)  # type: ignore
         hb = entropy(rhob)
 
     # pure system
     else:
         hab = 0.0
-        rhoa = reduced_density_matrix(s, traceout)
+        rhoa = reduced_density_matrix(s, traceout, d=d)
         ha = hb = entropy(rhoa)
 
     return ha + hb - hab
@@ -2591,7 +2642,7 @@ def mutual_information(s: Tensor, cut: Union[int, List[int]]) -> Tensor:
 # measurement results and transformations and correlations below
 
 
-def count_s2d(srepr: Tuple[Tensor, Tensor], n: int) -> Tensor:
+def count_s2d(srepr: Tuple[Tensor, Tensor], n: int, d: Optional[int] = None) -> Tensor:
     """
     measurement shots results, sparse tuple representation to dense representation
     count_vector to count_tuple
@@ -2600,11 +2651,14 @@ def count_s2d(srepr: Tuple[Tensor, Tensor], n: int) -> Tensor:
     :type srepr: Tuple[Tensor, Tensor]
     :param n: number of qubits
     :type n: int
+    :param d: [description], defaults to None
+    :type d: int, optional
     :return: [description]
     :rtype: Tensor
     """
+    d = 2 if d is None else d
     return backend.scatter(
-        backend.cast(backend.zeros([2**n]), srepr[1].dtype),
+        backend.cast(backend.zeros([d**n]), srepr[1].dtype),
         backend.reshape(srepr[0], [-1, 1]),
         srepr[1],
     )
@@ -2647,25 +2701,36 @@ def count_d2s(drepr: Tensor, eps: float = 1e-7) -> Tuple[Tensor, Tensor]:
 count_t2v = count_d2s
 
 
-def sample_int2bin(sample: Tensor, n: int) -> Tensor:
+def sample_int2bin(sample: Tensor, n: int, d: Optional[int] = None) -> Tensor:
     """
-    int sample to bin sample
+    Convert linear-index samples to per-site digits (base-d).
 
-    :param sample: in shape [trials] of int elements in the range [0, 2**n)
+    :param sample: shape [trials], integers in [0, d**n)
     :type sample: Tensor
-    :param n: number of qubits
+    :param n: number of sites
     :type n: int
-    :return: in shape [trials, n] of element (0, 1)
+    :param d: local dimension, defaults to 2
+    :type d: int, optional
+    :return: shape [trials, n], entries in [0, d-1]
     :rtype: Tensor
     """
-    confg = backend.mod(
-        backend.right_shift(sample[..., None], backend.reverse(backend.arange(n))),
-        2,
-    )
-    return confg
+    d = 2 if d is None else d
+    if d == 2:
+        return backend.mod(
+            backend.right_shift(sample[..., None], backend.reverse(backend.arange(n))),
+            2,
+        )
+    else:
+        pos = backend.reverse(backend.arange(n))
+        base = backend.power(d, pos)
+        digits = backend.mod(
+            backend.floor(backend.divide(sample[..., None], base)),  # ⌊sample / d**pos⌋
+            d,
+        )
+        return backend.cast(digits, "int32")
 
 
-def sample_bin2int(sample: Tensor, n: int) -> Tensor:
+def sample_bin2int(sample: Tensor, n: int, d: Optional[int] = None) -> Tensor:
     """
     bin sample to int sample
 
@@ -2676,88 +2741,98 @@ def sample_bin2int(sample: Tensor, n: int) -> Tensor:
     :return: in shape [trials]
     :rtype: Tensor
     """
-    power = backend.convert_to_tensor([2**j for j in reversed(range(n))])
+    d = 2 if d is None else d
+    power = backend.convert_to_tensor([d**j for j in reversed(range(n))])
     return backend.sum(sample * power, axis=-1)
 
 
 def sample2count(
-    sample: Tensor, n: int, jittable: bool = True
+    sample: Tensor,
+    n: int,
+    jittable: bool = True,
+    d: Optional[int] = None,
 ) -> Tuple[Tensor, Tensor]:
     """
-    sample_int to count_tuple
+    sample_int to count_tuple (indices, counts), size = d**n
 
-    :param sample: _description_
-    :type sample: Tensor
-    :param n: _description_
-    :type n: int
-    :param jittable: _description_, defaults to True
-    :type jittable: bool, optional
-    :return: _description_
-    :rtype: Tuple[Tensor, Tensor]
+    :param sample: linear-index samples, shape [shots]
+    :param n: number of sites
+    :param jittable: whether to return fixed-size outputs (backend dependent)
+    :param d: local dimension per site, default 2 (qubit)
+    :return: (unique_indices, counts)
     """
-    d = 2**n
+    d = 2 if d is None else d
+    size = d**n
     if not jittable:
         results = backend.unique_with_counts(sample)  # non-jittable
-    else:  # jax specified
-        results = backend.unique_with_counts(sample, size=d, fill_value=-1)
+    else:  # jax specified / fixed-size
+        results = backend.unique_with_counts(sample, size=size, fill_value=-1)
     return results
 
 
-def count_vector2dict(count: Tensor, n: int, key: str = "bin") -> Dict[Any, int]:
+def count_vector2dict(
+    count: Tensor, n: int, key: str = "bin", d: Optional[int] = None
+) -> Dict[Any, int]:
     """
-    convert_vector to count_dict_bin or count_dict_int
+    Convert count_vector to count_dict_bin or count_dict_int.
+    For d>10 cases, a base-d string (0-9A-Z) is used.
 
-    :param count: tensor in shape [2**n]
+    :param count: tensor in shape [d**n]
     :type count: Tensor
-    :param n: number of qubits
+    :param n: number of sites
     :type n: int
     :param key: can be "int" or "bin", defaults to "bin"
     :type key: str, optional
-    :return: _description_
-    :rtype: _type_
+    :param d: local dimension (default 2)
+    :type d: int, optional
+    :return: mapping from configuration to count
+    :rtype: Dict[Any, int]
     """
     from .interfaces import which_backend
-
+    d = 2 if d is None else d
     b = which_backend(count)
-    d = {i: b.numpy(count[i]).item() for i in range(2**n)}
+    out_int = {i: b.numpy(count[i]).item() for i in range(d**n)}
     if key == "int":
-        return d
+        return out_int
     else:
-        dn = {}
-        for k, v in d.items():
-            kn = str(bin(k))[2:].zfill(n)
-            dn[kn] = v
-        return dn
+        out_str = {}
+        for k, v in out_int.items():
+            kn = np.base_repr(k, base=d).zfill(n)
+            out_str[kn] = v
+        return out_str
 
 
 def count_tuple2dict(
-    count: Tuple[Tensor, Tensor], n: int, key: str = "bin"
+    count: Tuple[Tensor, Tensor], n: int, key: str = "bin", d: Optional[int] = None
 ) -> Dict[Any, int]:
     """
     count_tuple to count_dict_bin or count_dict_int
 
-    :param count: count_tuple format
+    :param count: count_tuple format (indices, counts)
     :type count: Tuple[Tensor, Tensor]
-    :param n: number of qubits
+    :param n: number of sites (qubits or qudits)
     :type n: int
     :param key: can be "int" or "bin", defaults to "bin"
     :type key: str, optional
+    :param d: local dimension, defaults to 2
+    :type d: int, optional
     :return: count_dict
-    :rtype: _type_
+    :rtype: Dict[Any, int]
     """
-    d = {
+    d = 2 if d is None else d
+    out_int = {
         backend.numpy(i).item(): backend.numpy(j).item()
         for i, j in zip(count[0], count[1])
         if i >= 0
     }
     if key == "int":
-        return d
+        return out_int
     else:
-        dn = {}
-        for k, v in d.items():
-            kn = str(bin(k))[2:].zfill(n)
-            dn[kn] = v
-        return dn
+        out_str = {}
+        for k, v in out_int.items():
+            kn = np.base_repr(k, base=d).zfill(n)
+            out_str[kn] = v
+        return out_str
 
 
 @partial(arg_alias, alias_dict={"counts": ["shots"], "format": ["format_"]})
@@ -2769,8 +2844,9 @@ def measurement_counts(
     random_generator: Optional[Any] = None,
     status: Optional[Tensor] = None,
     jittable: bool = False,
+    d: Optional[int] = None,
 ) -> Any:
-    """
+    r"""
     Simulate the measuring of each qubit of ``p`` in the computational basis,
     thus producing output like that of ``qiskit``.
 
@@ -2785,6 +2861,7 @@ def measurement_counts(
     "count_tuple": # (np.array([0]), np.array([2]))
 
     "count_dict_bin": # {"00": 2, "01": 0, "10": 0, "11": 0}
+    / for cases d\in [10, 36], "10" -> "A", ..., "35" -> "Z"
 
     "count_dict_int": # {0: 2, 1: 0, 2: 0, 3: 0}
 
@@ -2836,21 +2913,22 @@ def measurement_counts(
             state /= backend.norm(state)
             pi = backend.real(backend.conj(state) * state)
         pi = backend.reshape(pi, [-1])
-    d = int(backend.shape_tuple(pi)[0])
-    n = int(np.log(d) / np.log(2) + 1e-8)
+
+    local_d = 2 if d is None else d
+    total_dim = int(backend.shape_tuple(pi)[0])
+    n = _infer_num_sites(total_dim, local_d)
+
     if (counts is None) or counts <= 0:
         if format == "count_vector":
             return pi
         elif format == "count_tuple":
             return count_d2s(pi)
         elif format == "count_dict_bin":
-            return count_vector2dict(pi, n, key="bin")
+            return count_vector2dict(pi, n, key="bin", d=local_d)
         elif format == "count_dict_int":
-            return count_vector2dict(pi, n, key="int")
+            return count_vector2dict(pi, n, key="int", d=local_d)
         else:
-            raise ValueError(
-                "unsupported format %s for analytical measurement" % format
-            )
+            raise ValueError(f"unsupported format {format} for analytical measurement")
     else:
         raw_counts = backend.probability_sample(
             counts, pi, status=status, g=random_generator
@@ -2861,7 +2939,7 @@ def measurement_counts(
         # raw_counts = backend.stateful_randc(
         # random_generator, a=drange, shape=counts, p=pi
         # )
-        return sample2all(raw_counts, n, format=format, jittable=jittable)
+        return sample2all(raw_counts, n, format=format, jittable=jittable, d=local_d)
 
 
 measurement_results = measurement_counts
@@ -2869,52 +2947,62 @@ def measurement_counts(
 
 @partial(arg_alias, alias_dict={"format": ["format_"]})
 def sample2all(
-    sample: Tensor, n: int, format: str = "count_vector", jittable: bool = False
+    sample: Tensor,
+    n: int,
+    format: str = "count_vector",
+    jittable: bool = False,
+    d: Optional[int] = None,
 ) -> Any:
     """
-    transform ``sample_int`` or ``sample_bin`` form results to other forms specified by ``format``
+    transform ``sample_int`` or ``sample_bin`` results to other forms specified by ``format``
 
-    :param sample: measurement shots results in ``sample_int`` or ``sample_bin`` format
+    :param sample: measurement shots results in ``sample_int`` (shape [shots]) or ``sample_bin`` (shape [shots, n])
     :type sample: Tensor
-    :param n: number of qubits
+    :param n: number of sites
     :type n: int
-    :param format: see the doc in the doc in :py:meth:`tensorcircuit.quantum.measurement_results`,
-        defaults to "count_vector"
+    :param format: see :py:meth:`tensorcircuit.quantum.measurement_results`, defaults to "count_vector"
     :type format: str, optional
     :param jittable: only applicable to count transformation in jax backend, defaults to False
     :type jittable: bool, optional
+    :param d: local dimension (2 for qubit; >2 for qudit), defaults to 2
+    :type d: Optional[int]
     :return: measurement results specified as ``format``
     :rtype: Any
     """
+    d = 2 if d is None else int(d)
+
     if len(backend.shape_tuple(sample)) == 1:
         sample_int = sample
-        sample_bin = sample_int2bin(sample, n)
+        sample_bin = sample_int2bin(sample, n, d=d)
     elif len(backend.shape_tuple(sample)) == 2:
-        sample_int = sample_bin2int(sample, n)
+        sample_int = sample_bin2int(sample, n, d=d)
         sample_bin = sample
     else:
         raise ValueError("unrecognized tensor shape for sample")
+
     if format == "sample_int":
         return sample_int
     elif format == "sample_bin":
         return sample_bin
     else:
-        count_tuple = sample2count(sample_int, n, jittable)
+        count_tuple = sample2count(sample_int, n, jittable=jittable, d=d)
         if format == "count_tuple":
             return count_tuple
         elif format == "count_vector":
-            return count_s2d(count_tuple, n)
+            return count_s2d(count_tuple, n, d=d)
         elif format == "count_dict_bin":
-            return count_tuple2dict(count_tuple, n, key="bin")
+            return count_tuple2dict(count_tuple, n, key="bin", d=d)
         elif format == "count_dict_int":
-            return count_tuple2dict(count_tuple, n, key="int")
+            return count_tuple2dict(count_tuple, n, key="int", d=d)
         else:
             raise ValueError(
-                "unsupported format %s for finite shots measurement" % format
+                f"unsupported format {format} for finite shots measurement"
             )
 
 
-def spin_by_basis(n: int, m: int, elements: Tuple[int, int] = (1, -1)) -> Tensor:
+def spin_by_basis(
+    n: int, m: int, elements: Tuple[int, int] = (1, -1), d: Optional[int] = None
+) -> Tensor:
     """
     Generate all n-bitstrings as an array, each row is a bitstring basis.
     Return m-th col.
@@ -2934,67 +3022,109 @@ def spin_by_basis(n: int, m: int, elements: Tuple[int, int] = (1, -1)) -> Tensor
         all bitstring basis.
     :rtype: Tensor
     """
-    s = backend.tile(
-        backend.cast(
-            backend.convert_to_tensor(np.array([[elements[0]], [elements[1]]])), "int32"
-        ),
-        [2**m, int(2 ** (n - m - 1))],
-    )
+    d = len(elements) if d is None else d
+
+    col = backend.convert_to_tensor(np.array(elements, dtype=np.int32).reshape(-1, 1))
+    s = backend.tile(backend.cast(col, "int32"), [d**m, int(d ** (n - m - 1))])
     return backend.reshape(s, [-1])
 
 
-def correlation_from_samples(index: Sequence[int], results: Tensor, n: int) -> Tensor:
+def correlation_from_samples(
+    index: Sequence[int],
+    results: Tensor,
+    n: int,
+    d: int = 2,
+    elements: Optional[Sequence[float]] = None,
+) -> Tensor:
     r"""
-    Compute :math:`\prod_{i\in \\text{index}} s_i (s=\pm 1)`,
-    Results is in the format of "sample_int" or "sample_bin"
+    Compute :math:`\prod_{i\in \text{index}} s_i` from measurement shots,
+    where each site value :math:`s_i` is mapped from the digit outcome.
 
-    :param index: list of int, indicating the position in the bitstring
-    :type index: Sequence[int]
-    :param results: sample tensor
-    :type results: Tensor
-    :param n: number of qubits
-    :type n: int
-    :return: Correlation expectation from measurement shots
-    :rtype: Tensor
+    Results can be "sample_int" ([shots]) or "sample_bin" ([shots, n]).
+
+    :param index: positions in the basis string
+    :param results: samples tensor
+    :param n: number of sites
+    :param d: local dimension (default 2)
+    :param elements: optional mapping of length d from outcome {0..d-1} to values s.
+                     If None and d==2, defaults to (1, -1) via the original formula.
+    :return: correlation estimate (mean over shots)
     """
     if len(backend.shape_tuple(results)) == 1:
-        results = sample_int2bin(results, n)
-    results = 1 - results * 2
-    r = results[:, index[0]]
+        results = sample_int2bin(results, n, d=d)
+
+    if d == 2 and elements is None:
+        svals = 1 - results * 2  # 0->+1, 1->-1
+        r = svals[:, index[0]]
+        for i in index[1:]:
+            r *= svals[:, i]
+        r = backend.cast(r, rdtypestr)
+        return backend.mean(r)
+
+    if elements is None:
+        raise ValueError(
+            f"correlation_from_samples requires `elements` mapping for d={d}; "
+            f"e.g., for qutrit you might pass elements=(1.0,0.0,-1.0)."
+        )
+    if len(elements) != d:
+        raise ValueError(f"`elements` length {len(elements)} != d={d}")
+
+    evec = backend.cast(backend.convert_to_tensor(np.asarray(elements)), rdtypestr)
+
+    col = backend.cast(results[:, index[0]], "int32")
+    r = backend.gather1d(evec, col)  # shape [shots]
+
     for i in index[1:]:
-        r *= results[:, i]
+        col = backend.cast(results[:, i], "int32")
+        r *= backend.gather1d(evec, col)
+
     r = backend.cast(r, rdtypestr)
     return backend.mean(r)
 
 
-def correlation_from_counts(index: Sequence[int], results: Tensor) -> Tensor:
+def correlation_from_counts(
+    index: Sequence[int],
+    results: Tensor,
+    d: Optional[int] = None,
+    elements: Optional[Sequence[float]] = None,
+) -> Tensor:
     r"""
-    Compute :math:`\prod_{i\in \\text{index}} s_i`,
-    where the probability for each bitstring is given as a vector ``results``.
-    Results is in the format of "count_vector"
+    Compute :math:`\prod_{i\in \text{index}} s_i` where the probability for each
+    basis label is given by a count/probability vector ``results`` ("count_vector").
 
-    :Example:
-
-    >>> prob = tc.array_to_tensor(np.array([0.6, 0.4, 0, 0]))
-    >>> qu.correlation_from_counts([0, 1], prob)
-    (0.20000002+0j)
-    >>> qu.correlation_from_counts([1], prob)
-    (0.20000002+0j)
-
-    :param index: list of int, indicating the position in the bitstring
-    :type index: Sequence[int]
-    :param results: probability vector of shape 2^n
-    :type results: Tensor
-    :return: Correlation expectation from measurement shots.
-    :rtype: Tensor
+    :param index: positions in the basis string
+    :param results: probability/count vector of shape d**n (will be normalized)
+    :param d: local dimension (default 2)
+    :param elements: optional mapping of length d from digit {0..d-1} to values s.
+                     If None and d==2, defaults to (1, -1). For d>2, must be provided.
+    :return: correlation expectation from counts
     """
+    d = 2 if d is None else int(d)
+    if d != 2:
+        raise NotImplementedError(f"`d={d}` not implemented.")
+
     results = backend.reshape(results, [-1])
     results = backend.cast(results, rdtypestr)
     results /= backend.sum(results)
-    n = int(np.log(results.shape[0]) / np.log(2))
+
+    n = _infer_num_sites(int(results.shape[0]), d=d)
+
+    if d == 2 and elements is None:
+        elems = (1, -1)
+    else:
+        if elements is None or len(elements) != d:
+            raise ValueError(
+                f"`elements` must be provided with length d={d} for qudit; got {elements}."
+            )
+        elems = tuple(elements)  # type: ignore
+
+    acc = results
     for i in index:
-        results = results * backend.cast(spin_by_basis(n, i), results.dtype)
-    return backend.sum(results)
+        acc = acc * backend.cast(
+            spin_by_basis(n, int(i), elements=elems, d=d), acc.dtype
+        )
+
+    return backend.sum(acc)
 
 
 # @op2tensor

From d92d6e160e6038fefed53a55dfb4d24a45b9ca39 Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Wed, 27 Aug 2025 10:19:59 +0800
Subject: [PATCH 05/55] black formatted file.

---
 tensorcircuit/quantum.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorcircuit/quantum.py b/tensorcircuit/quantum.py
index a6b60a66..3454e27a 100644
--- a/tensorcircuit/quantum.py
+++ b/tensorcircuit/quantum.py
@@ -2789,6 +2789,7 @@ def count_vector2dict(
     :rtype: Dict[Any, int]
     """
     from .interfaces import which_backend
+
     d = 2 if d is None else d
     b = which_backend(count)
     out_int = {i: b.numpy(count[i]).item() for i in range(d**n)}

From 01d2538c5d4e3ad6dde67f33b347ca8deae6f95f Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Wed, 27 Aug 2025 10:21:23 +0800
Subject: [PATCH 06/55] Expanded the any_gate function, which is main entry
 function for the qudit gates. And optimized some code writing.

---
 tensorcircuit/gates.py | 50 +++++++++++++++++++++++-------------------
 1 file changed, 28 insertions(+), 22 deletions(-)

diff --git a/tensorcircuit/gates.py b/tensorcircuit/gates.py
index 459666ab..e9f80ae3 100644
--- a/tensorcircuit/gates.py
+++ b/tensorcircuit/gates.py
@@ -34,6 +34,12 @@
 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 @@ def num_to_tensor(*num: Union[float, Tensor], dtype: Optional[str] = None) -> An
 
 
 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 @@ def f(*args: Any, **kws: Any) -> Any:
 
             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 @@ def f(*args: Any, **kws: Any) -> Any:
             # 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 @@ def __init__(
         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)

From 2a9d8ce075ec666499e439b011ccfd52d16cbdef Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Wed, 27 Aug 2025 10:29:41 +0800
Subject: [PATCH 07/55] Add a parameter _d for abstractcircuit, which is 2 in
 Default.

---
 tensorcircuit/abstractcircuit.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorcircuit/abstractcircuit.py b/tensorcircuit/abstractcircuit.py
index 3f3d5951..6dd6d07d 100644
--- a/tensorcircuit/abstractcircuit.py
+++ b/tensorcircuit/abstractcircuit.py
@@ -68,6 +68,7 @@
 
 class AbstractCircuit:
     _nqubits: int
+    _d: int = 2
     _qir: List[Dict[str, Any]]
     _extra_qir: List[Dict[str, Any]]
     inputs: Tensor

From 9baccfb0627a7e222ac93dd3d1e0d4acb430b324 Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Wed, 27 Aug 2025 10:32:26 +0800
Subject: [PATCH 08/55] Removed hard code in circuit.py (correspond to only
 qubit system), Circuit is now suitable to do general calculation for
 d-dimensional systems.

---
 tensorcircuit/circuit.py | 112 +++++++++++++++++++++------------------
 1 file changed, 61 insertions(+), 51 deletions(-)

diff --git a/tensorcircuit/circuit.py b/tensorcircuit/circuit.py
index b18bdf49..01a53e08 100644
--- a/tensorcircuit/circuit.py
+++ b/tensorcircuit/circuit.py
@@ -1,5 +1,7 @@
 """
-Quantum circuit: the state simulator
+Quantum circuit: the state simulator.
+Supports qubit (d=2) and qudit (3 <= d <= 36) systems.
+ For string-encoded samples/counts, digits use 0–9A–Z where A=10, …, Z=35.
 """
 
 # pylint: disable=invalid-name
@@ -13,7 +15,7 @@
 
 from . import gates
 from . import channels
-from .cons import backend, contractor, dtypestr, npdtype
+from .cons import backend, contractor, dtypestr, npdtype, _ALPHABET
 from .quantum import QuOperator, identity
 from .simplify import _full_light_cone_cancel
 from .basecircuit import BaseCircuit
@@ -23,7 +25,7 @@
 
 
 class Circuit(BaseCircuit):
-    """
+    r"""
     ``Circuit`` class.
     Simple usage demo below.
 
@@ -42,17 +44,21 @@ class Circuit(BaseCircuit):
     def __init__(
         self,
         nqubits: int,
+        dim: Optional[int] = None,
         inputs: Optional[Tensor] = None,
         mps_inputs: Optional[QuOperator] = None,
         split: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
     ) -> None:
-        """
+        r"""
         Circuit object based on state simulator.
 
         :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 inputs: If not None, the initial state of the circuit is taken as ``inputs``
-            instead of :math:`\\vert 0\\rangle^n` qubits, defaults to None.
+            instead of :math:`\vert 0 \rangle^n` qubits, defaults to None.
         :type inputs: Optional[Tensor], optional
         :param mps_inputs: QuVector for a MPS like initial wavefunction.
         :type mps_inputs: Optional[QuOperator]
@@ -60,6 +66,13 @@ def __init__(
             ``max_singular_values`` and ``max_truncation_err``.
         :type split: Optional[Dict[str, Any]]
         """
+        self._d = 2 if dim is None else dim
+        if not kwargs.get("qudit", False) and self._d != 2:
+            raise ValueError(
+                f"Circuit only supports qubits (dim=2). "
+                f"You passed dim={self._d}. Please use `QuditCircuit` instead."
+            )
+
         self.inputs = inputs
         self.mps_inputs = mps_inputs
         self.split = split
@@ -67,21 +80,22 @@ def __init__(
 
         self.circuit_param = {
             "nqubits": nqubits,
+            "dim": dim,
             "inputs": inputs,
             "mps_inputs": mps_inputs,
             "split": split,
         }
         if (inputs is None) and (mps_inputs is None):
-            nodes = self.all_zero_nodes(nqubits)
+            nodes = self.all_zero_nodes(nqubits, d=self._d)
             self._front = [n.get_edge(0) for n in nodes]
         elif inputs is not None:  # provide input function
             inputs = backend.convert_to_tensor(inputs)
             inputs = backend.cast(inputs, dtype=dtypestr)
             inputs = backend.reshape(inputs, [-1])
             N = inputs.shape[0]
-            n = int(np.log(N) / np.log(2))
+            n = int(np.log(N) / np.log(self._d))
             assert n == nqubits or n == 2 * nqubits
-            inputs = backend.reshape(inputs, [2 for _ in range(n)])
+            inputs = backend.reshape(inputs, [self._d for _ in range(n)])
             inputs = Gate(inputs)
             nodes = [inputs]
             self._front = [inputs.get_edge(i) for i in range(n)]
@@ -178,27 +192,14 @@ def mid_measurement(self, index: int, keep: int = 0) -> Tensor:
 
         :param index: The index of qubit that the Z direction postselection applied on.
         :type index: int
-        :param keep: 0 for spin up, 1 for spin down, defaults to be 0.
+        :param keep: the post-selected digit in {0, ..., d-1}, defaults to be 0.
         :type keep: int, optional
         """
         # normalization not guaranteed
-        # assert keep in [0, 1]
-        if keep < 0.5:
-            gate = np.array(
-                [
-                    [1.0],
-                    [0.0],
-                ],
-                dtype=npdtype,
-            )
-        else:
-            gate = np.array(
-                [
-                    [0.0],
-                    [1.0],
-                ],
-                dtype=npdtype,
-            )
+        gate = np.array(
+            [[0.0] if _idx != keep else [1.0] for _idx in range(self._d)],
+            dtype=npdtype,
+        )
 
         mg1 = tn.Node(gate)
         mg2 = tn.Node(gate)
@@ -479,7 +480,7 @@ def step_function(x: Tensor) -> Tensor:
         if get_gate_from_index is None:
             raise ValueError("no `get_gate_from_index` implementation is provided")
         g = get_gate_from_index(r, kraus)
-        g = backend.reshape(g, [2 for _ in range(sites * 2)])
+        g = backend.reshape(g, [self._d for _ in range(sites * 2)])
         self.any(*index, unitary=g, name=name)  # type: ignore
         return r
 
@@ -680,7 +681,7 @@ def _meta_apply_channels(cls) -> None:
             Apply %s quantum channel on the circuit.
             See :py:meth:`tensorcircuit.channels.%schannel`
 
-            :param index: Qubit number that the gate applies on.
+            :param index: Site index that the gate applies on.
             :type index: int.
             :param status: uniform external random number between 0 and 1
             :type status: Tensor
@@ -737,8 +738,8 @@ def get_quoperator(self) -> QuOperator:
         :return: ``QuOperator`` object for the circuit unitary (open indices for the input state)
         :rtype: QuOperator
         """
-        mps = identity([2 for _ in range(self._nqubits)])
-        c = Circuit(self._nqubits)
+        mps = identity([self._d for _ in range(self._nqubits)])
+        c = Circuit(self._nqubits, self._d)
         ns, es = self._copy()
         c._nodes = ns
         c._front = es
@@ -758,8 +759,8 @@ def matrix(self) -> Tensor:
         :return: The circuit unitary matrix
         :rtype: Tensor
         """
-        mps = identity([2 for _ in range(self._nqubits)])
-        c = Circuit(self._nqubits)
+        mps = identity([self._d for _ in range(self._nqubits)])
+        c = Circuit(self._nqubits, self._d)
         ns, es = self._copy()
         c._nodes = ns
         c._front = es
@@ -772,6 +773,9 @@ def measure_reference(
         """
         Take measurement on the given quantum lines by ``index``.
 
+        Return format:
+        - For d <= 36, the sample is a base-d string using 0–9A–Z (A=10,…).
+
         :Example:
 
         >>> c = tc.Circuit(3)
@@ -800,10 +804,9 @@ def measure_reference(
                 if i != j:
                     e ^ edge2[i]
             for i in range(len(sample)):
-                if sample[i] == "0":
-                    m = np.array([1, 0], dtype=npdtype)
-                else:
-                    m = np.array([0, 1], dtype=npdtype)
+                m = np.array([0 for _ in range(self._d)], dtype=npdtype)
+                m[int(sample[i])] = 1
+
                 nodes1.append(tn.Node(m))
                 nodes1[-1].get_edge(0) ^ edge1[index[i]]
                 nodes2.append(tn.Node(m))
@@ -814,15 +817,13 @@ def measure_reference(
                 / p
                 * contractor(nodes1, output_edge_order=[edge1[j], edge2[j]]).tensor
             )
-            pu = rho[0, 0]
-            r = backend.random_uniform([])
-            r = backend.real(backend.cast(r, dtypestr))
-            if r < backend.real(pu):
-                sample += "0"
-                p = p * pu
-            else:
-                sample += "1"
-                p = p * (1 - pu)
+            probs = backend.real(backend.diagonal(rho))
+            probs /= backend.sum(probs)
+            outcome = np.random.choice(self._d, p=probs)
+
+            sample += _ALPHABET[outcome]
+            p *= float(probs[outcome])
+
         if with_prob:
             return sample, p
         else:
@@ -842,6 +843,10 @@ def expectation(
     ) -> Tensor:
         """
         Compute the expectation of corresponding operators.
+        For qudit (d > 2),
+        ensure that operator tensor shapes are consistent with d (each site contributes two axes of size d).
+
+        Noise shorthand (via noise_conf) is qubit-only; for d>2, use explicit operators.
 
         :Example:
 
@@ -883,8 +888,6 @@ def expectation(
         :return: Tensor with one element
         :rtype: Tensor
         """
-        from .noisemodel import expectation_noisfy
-
         if noise_conf is None:
             # if not reuse:
             #     nodes1, edge1 = self._copy()
@@ -899,6 +902,8 @@ def expectation(
                 nodes1 = _full_light_cone_cancel(nodes1)
             return contractor(nodes1).tensor
         else:
+            from .noisemodel import expectation_noisfy
+
             return expectation_noisfy(
                 self,
                 *ops,
@@ -916,12 +921,14 @@ def expectation(
 def expectation(
     *ops: Tuple[tn.Node, List[int]],
     ket: Tensor,
+    d: Optional[int] = None,
     bra: Optional[Tensor] = None,
     conj: bool = True,
     normalization: bool = False,
 ) -> Tensor:
     """
     Compute :math:`\\langle bra\\vert ops \\vert ket\\rangle`.
+    For qudit systems (d>2), ops must be reshaped with per-site axes of length d.
 
     Example 1 (:math:`bra` is same as :math:`ket`)
 
@@ -966,6 +973,8 @@ def expectation(
     :type ket: Tensor
     :param bra: :math:`bra`, defaults to None, which is the same as ``ket``.
     :type bra: Optional[Tensor], optional
+    :param d: dimension of the circuit (defaults to 2)
+    :type d: int, optional
     :param conj: :math:`bra` changes to the adjoint matrix of :math:`bra`, defaults to True.
     :type conj: bool, optional
     :param normalization: Normalize the :math:`ket` and :math:`bra`, defaults to False.
@@ -974,6 +983,7 @@ def expectation(
     :return: The result of :math:`\\langle bra\\vert ops \\vert ket\\rangle`.
     :rtype: Tensor
     """
+    d = 2 if d is None else d
     if bra is None:
         bra = ket
     if isinstance(ket, QuOperator):
@@ -987,7 +997,7 @@ def expectation(
         for op, index in ops:
             if not isinstance(op, tn.Node):
                 # op is only a matrix
-                op = backend.reshape2(op)
+                op = backend.reshaped(op, d)
                 op = gates.Gate(op)
             if isinstance(index, int):
                 index = [index]
@@ -1011,8 +1021,8 @@ def expectation(
         if conj is True:
             bra = backend.conj(bra)
         ket = backend.reshape(ket, [-1])
-        ket = backend.reshape2(ket)
-        bra = backend.reshape2(bra)
+        ket = backend.reshaped(ket, d)
+        bra = backend.reshaped(bra, d)
         n = len(backend.shape_tuple(ket))
         ket = Gate(ket)
         bra = Gate(bra)
@@ -1024,7 +1034,7 @@ def expectation(
         for op, index in ops:
             if not isinstance(op, tn.Node):
                 # op is only a matrix
-                op = backend.reshape2(op)
+                op = backend.reshaped(op, d)
                 op = gates.Gate(op)
             if isinstance(index, int):
                 index = [index]

From e5ff8cab23da86e353f03b05ff5e99b7f491831c Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Wed, 27 Aug 2025 10:41:05 +0800
Subject: [PATCH 09/55] decode the labels for new representation of states,
 0-9A-Z.

---
 tensorcircuit/basecircuit.py | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/tensorcircuit/basecircuit.py b/tensorcircuit/basecircuit.py
index e39f0c6b..5b9ce73a 100644
--- a/tensorcircuit/basecircuit.py
+++ b/tensorcircuit/basecircuit.py
@@ -32,6 +32,30 @@
 Tensor = Any
 
 
+def _decode_basis_label(label: str, d: int, n: int) -> List[int]:
+    if d > 36:
+        raise NotImplementedError(
+            f"String basis label supports d<=36 (0–9A–Z). Got d={d}. "
+            "Use an integer array/tensor of length n instead."
+        )
+    s = label.upper()
+    if len(s) != n:
+        raise ValueError(f"Basis label length mismatch: expect {n}, got {len(s)}")
+    digits = []
+    for ch in s:
+        if ch not in _ALPHABET:
+            raise ValueError(
+                f"Invalid character '{ch}' in basis label (allowed 0–9A–Z)."
+            )
+        v = _ALPHABET.index(ch)
+        if v >= d:
+            raise ValueError(
+                f"Digit '{ch}' (= {v}) out of range for base-d with d={d}."
+            )
+        digits.append(v)
+    return digits
+
+
 class BaseCircuit(AbstractCircuit):
     _nodes: List[tn.Node]
     _front: List[tn.Edge]

From 5b7ffe47344e37e211a552ccd2c8b04d46361a51 Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Wed, 27 Aug 2025 10:41:48 +0800
Subject: [PATCH 10/55] change site label, when dim>2, it changes from `qb-` to
 `qd-`

---
 tensorcircuit/basecircuit.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorcircuit/basecircuit.py b/tensorcircuit/basecircuit.py
index 5b9ce73a..c0d889e9 100644
--- a/tensorcircuit/basecircuit.py
+++ b/tensorcircuit/basecircuit.py
@@ -66,6 +66,7 @@ class BaseCircuit(AbstractCircuit):
 
     @staticmethod
     def all_zero_nodes(n: int, d: int = 2, prefix: str = "qb-") -> List[tn.Node]:
+        prefix = "qd-" if d > 2 else prefix
         l = [0.0 for _ in range(d)]
         l[0] = 1.0
         nodes = [

From 10627e6fd835a97f9e002eeed6ee2daafa401d27 Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Wed, 27 Aug 2025 11:07:27 +0800
Subject: [PATCH 11/55] black .

---
 tensorcircuit/basecircuit.py | 224 ++++++++++++++++++++++++-----------
 1 file changed, 152 insertions(+), 72 deletions(-)

diff --git a/tensorcircuit/basecircuit.py b/tensorcircuit/basecircuit.py
index c0d889e9..c84acf6c 100644
--- a/tensorcircuit/basecircuit.py
+++ b/tensorcircuit/basecircuit.py
@@ -1,5 +1,9 @@
 """
 Quantum circuit: common methods for all circuit classes as MixIn
+
+Note:
+  - Supports qubit (d = 2) and qudit (d >= 2) systems.
+  - For string-encoded samples/counts when d <= 36, digits use base-d characters 0–9A–Z (A = 10, …, Z = 35).
 """
 
 # pylint: disable=invalid-name
@@ -21,9 +25,10 @@
     sample_int2bin,
     sample_bin2int,
     sample2all,
+    _infer_num_sites,
 )
 from .abstractcircuit import AbstractCircuit
-from .cons import npdtype, backend, dtypestr, contractor, rdtypestr
+from .cons import npdtype, backend, dtypestr, contractor, rdtypestr, _ALPHABET
 from .simplify import _split_two_qubit_gate
 from .utils import arg_alias
 
@@ -314,7 +319,7 @@ def expectation_before(
         for op, index in ops:
             if not isinstance(op, tn.Node):
                 # op is only a matrix
-                op = backend.reshape2(op)
+                op = backend.reshaped(op, d=self._d)
                 op = backend.cast(op, dtype=dtypestr)
                 op = gates.Gate(op)
             else:
@@ -380,12 +385,12 @@ def to_qir(self) -> List[Dict[str, Any]]:
 
     def perfect_sampling(self, status: Optional[Tensor] = None) -> Tuple[str, float]:
         """
-        Sampling bistrings from the circuit output based on quantum amplitudes.
+        Sampling base-d strings (0–9A–Z when d <= 36) from the circuit output based on quantum amplitudes.
         Reference: arXiv:1201.3974.
 
         :param status: external randomness, with shape [nqubits], defaults to None
         :type status: Optional[Tensor]
-        :return: Sampled bit string and the corresponding theoretical probability.
+        :return: Sampled base-d string and the corresponding theoretical probability.
         :rtype: Tuple[str, float]
         """
         return self.measure_jit(*range(self._nqubits), with_prob=True, status=status)
@@ -394,10 +399,10 @@ def measure_jit(
         self, *index: int, with_prob: bool = False, status: Optional[Tensor] = None
     ) -> Tuple[Tensor, Tensor]:
         """
-        Take measurement to the given quantum lines.
+        Take measurement on the given site indices (computational basis).
         This method is jittable is and about 100 times faster than unjit version!
 
-        :param index: Measure on which quantum line.
+        :param index: Measure on which site (wire) index.
         :type index: int
         :param with_prob: If true, theoretical probability is also returned.
         :type with_prob: bool, optional
@@ -408,9 +413,8 @@ def measure_jit(
         """
         # finally jit compatible ! and much faster than unjit version ! (100x)
         sample: List[Tensor] = []
-        p = 1.0
-        p = backend.convert_to_tensor(p)
-        p = backend.cast(p, dtype=rdtypestr)
+        one_r = backend.cast(backend.convert_to_tensor(1.0), rdtypestr)
+        p = one_r
         for k, j in enumerate(index):
             if self.is_dm is False:
                 nodes1, edge1 = self._copy()
@@ -425,40 +429,88 @@ def measure_jit(
                 if i != j:
                     e ^ edge2[i]
             for i in range(k):
-                m = (1 - sample[i]) * gates.array_to_tensor(np.array([1, 0])) + sample[
-                    i
-                ] * gates.array_to_tensor(np.array([0, 1]))
-                newnodes.append(Gate(m))
-                newnodes[-1].id = id(newnodes[-1])
-                newnodes[-1].is_dagger = False
-                newnodes[-1].flag = "measurement"
-                newnodes[-1].get_edge(0) ^ edge1[index[i]]
-                newnodes.append(Gate(m))
-                newnodes[-1].id = id(newnodes[-1])
-                newnodes[-1].is_dagger = True
-                newnodes[-1].flag = "measurement"
-                newnodes[-1].get_edge(0) ^ edge2[index[i]]
+                if self._d == 2:
+                    m = (1 - sample[i]) * gates.array_to_tensor(
+                        np.array([1, 0])
+                    ) + sample[i] * gates.array_to_tensor(np.array([0, 1]))
+                    g1 = Gate(m)
+                    g1.id = id(g1)
+                    g1.is_dagger = False
+                    g1.flag = "measurement"
+                    newnodes.append(g1)
+                    g1.get_edge(0) ^ edge1[index[i]]
+                    g2 = Gate(m)
+                    g2.id = id(g2)
+                    g2.is_dagger = True
+                    g2.flag = "measurement"
+                    newnodes.append(g2)
+                    g2.get_edge(0) ^ edge2[index[i]]
+                else:
+                    vec = backend.one_hot(backend.cast(sample[i], "int32"), self._d)
+                    v = backend.cast(vec, dtypestr)
+                    m = backend.tensordot(v, v, axes=0)
+                    g = Gate(m)
+                    g.id = id(g)
+                    g.is_dagger = False
+                    g.flag = "measurement"
+                    newnodes.append(g)
+                    g.get_edge(0) ^ edge1[index[i]]
+                    g.get_edge(1) ^ edge2[index[i]]
             rho = (
                 1
                 / backend.cast(p, dtypestr)
                 * contractor(newnodes, output_edge_order=[edge1[j], edge2[j]]).tensor
             )
-            pu = backend.real(rho[0, 0])
-            if status is None:
-                r = backend.implicit_randu()[0]
+            if self._d == 2:
+                pu = backend.real(rho[0, 0])
+                if status is None:
+                    r = backend.implicit_randu()[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)
             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)
-
-        sample = backend.stack(sample)
-        sample = backend.real(sample)
+                zero_r = backend.cast(backend.convert_to_tensor(0.0), rdtypestr)
+                tiny_r = backend.cast(backend.convert_to_tensor(1e-12), rdtypestr)
+                pu = backend.real(backend.diagonal(rho))
+                pu = backend.clip(pu, zero_r, one_r)
+                d = backend.shape_tuple(pu)[-1]
+                pu = pu + tiny_r * (
+                    backend.ones((d,), dtype=rdtypestr)
+                    / backend.cast(backend.convert_to_tensor(float(d)), rdtypestr)
+                )
+                pu = pu / backend.sum(pu)
+                cdf = backend.cumsum(pu)
+                if status is None:
+                    r = backend.implicit_randu()[0]
+                    r = backend.real(backend.cast(r, rdtypestr))
+                    phi = backend.cast(
+                        backend.convert_to_tensor(0.6180339887498948), rdtypestr
+                    )
+                    r = r + phi * backend.cast(
+                        backend.convert_to_tensor(k + 1), rdtypestr
+                    )
+                    r = r - backend.floor(r)
+                    r = backend.clip(r, zero_r, one_r - tiny_r)
+                else:
+                    r = backend.real(backend.cast(status[k], rdtypestr))
+                k_out = backend.searchsorted(cdf, r, side="right")
+                k_out = backend.clip(
+                    k_out,
+                    backend.cast(backend.convert_to_tensor(0), "int32"),
+                    backend.cast(backend.convert_to_tensor(d - 1), "int32"),
+                )
+                sample.append(backend.cast(k_out, rdtypestr))
+                p = p * backend.cast(pu[k_out], rdtypestr)
+        sample = backend.real(backend.stack(sample))
         if with_prob:
             return sample, p
         else:
@@ -468,35 +520,56 @@ def measure_jit(
 
     def amplitude_before(self, l: Union[str, Tensor]) -> List[Gate]:
         r"""
-        Returns the tensornetwor nodes for the amplitude of the circuit given the bitstring l.
-        For state simulator, it computes :math:`\langle l\vert \psi\rangle`,
-        for density matrix simulator, it computes :math:`Tr(\rho \vert l\rangle \langle 1\vert)`
+        Returns the tensornetwor nodes for the amplitude of the circuit given a computational-basis label ``l``.
+        For a state simulator, it computes :math:`\langle l \vert \psi\rangle`;
+        for a density-matrix simulator, it computes :math:`\mathrm{Tr}(\rho \vert l\rangle\langle l\vert)`.
         Note how these two are different up to a square operation.
 
-        :param l: The bitstring of 0 and 1s.
+        :Example:
+
+        >>> c = tc.Circuit(2)
+        >>> c.X(0)
+        >>> c.amplitude("10")  # d=2, per-qubit digits
+        array(1.+0.j, dtype=complex64)
+        >>> c.CNOT(0, 1)
+        >>> c.amplitude("11")
+        array(1.+0.j, dtype=complex64)
+
+        For qudits (d>2, d<=36):
+        >>> c = tc.Circuit(3, dim=12)
+        >>> c.amplitude("0A2")  # base-12 string, A stands for 10
+
+        :param l: Basis label.
+            - If a string: it must be a base-d string of length ``nqubits``, using 0–9A–Z (A=10,…,Z=35) when ``d<=36``.
+            - If a tensor/array/list: it should contain per-site integers in ``[0, d-1]`` with length ``nqubits``.
         :type l: Union[str, Tensor]
         :return: The tensornetwork nodes for the amplitude of the circuit.
         :rtype: List[Gate]
         """
+
+        def _basis_nod(_k: int) -> Tensor:
+            _vec = np.zeros((self._d,), dtype=npdtype)
+            _vec[_k] = 1.0
+            return _vec
+
         no, d_edges = self._copy()
         ms = []
         if self.is_dm:
             msconj = []
         if isinstance(l, str):
-            for s in l:
-                if s == "1":
-                    endn = np.array([0, 1], dtype=npdtype)
-                elif s == "0":
-                    endn = np.array([1, 0], dtype=npdtype)
-                ms.append(tn.Node(endn))
+            symbols = _decode_basis_label(l, d=self._d, n=self._nqubits)
+            for k in symbols:
+                n = _basis_nod(k)
+                ms.append(tn.Node(n))
                 if self.is_dm:
-                    msconj.append(tn.Node(endn))
-        else:  # l is Tensor
+                    msconj.append(tn.Node(n))
+        else:
             l = backend.cast(l, dtype=dtypestr)
             for i in range(self._nqubits):
-                endn = l[i] * gates.array_to_tensor(np.array([0, 1])) + (
-                    1 - l[i]
-                ) * gates.array_to_tensor(np.array([1, 0]))
+                endn = backend.cast(
+                    backend.one_hot(backend.cast(l[i], "int32"), self._d),
+                    dtype=dtypestr,
+                )
                 ms.append(tn.Node(endn))
                 if self.is_dm:
                     msconj.append(tn.Node(endn))
@@ -547,17 +620,18 @@ def amplitude(self, l: Union[str, Tensor]) -> Tensor:
 
     def probability(self) -> Tensor:
         """
-        get the 2^n length probability vector over computational basis
+        get the d^n length probability vector over computational basis
 
-        :return: probability vector
+        :return: probability vector of shape [d**n]
         :rtype: Tensor
         """
         s = self.state()  # type: ignore
         if self.is_dm is False:
-            p = backend.abs(s) ** 2
-
+            amp = backend.reshape(s, [-1])
+            p = backend.real(backend.abs(amp) ** 2)
         else:
-            p = backend.abs(backend.diagonal(s))
+            diag = backend.diagonal(s)
+            p = backend.real(backend.reshape(diag, [-1]))
         return p
 
     @partial(arg_alias, alias_dict={"format": ["format_"]})
@@ -571,7 +645,7 @@ def sample(
         status: Optional[Tensor] = None,
         jittable: bool = True,
     ) -> Any:
-        """
+        r"""
         batched sampling from state or circuit tensor network directly
 
         :param batch: number of samples, defaults to None
@@ -594,6 +668,7 @@ def sample(
                 "count_tuple": # (np.array([0]), np.array([2]))
 
                 "count_dict_bin": # {"00": 2, "01": 0, "10": 0, "11": 0}
+                    for cases d\in [11, 36], use 0–9A–Z digits (e.g., 'A' -> 10, …, 'Z' -> 35);
 
                 "count_dict_int": # {0: 2, 1: 0, 2: 0, 3: 0}
 
@@ -645,7 +720,7 @@ def perfect_sampling(key: Any) -> Any:
                 return r
             r = backend.stack([ri[0] for ri in r])  # type: ignore
             r = backend.cast(r, "int32")
-            ch = sample_bin2int(r, self._nqubits)
+            ch = sample_bin2int(r, self._nqubits, d=self._d)
         else:  # allow_state
             if batch is None:
                 nbatch = 1
@@ -670,7 +745,7 @@ def perfect_sampling(key: Any) -> Any:
             #     2,
             # )
             if format is None:  # for backward compatibility
-                confg = sample_int2bin(ch, self._nqubits)
+                confg = sample_int2bin(ch, self._nqubits, d=self._d)
                 prob = backend.gather1d(p, ch)
                 r = list(zip(confg, prob))  # type: ignore
                 if batch is None:
@@ -678,7 +753,9 @@ def perfect_sampling(key: Any) -> Any:
                 return r
         if self._nqubits > 35:
             jittable = False
-        return sample2all(sample=ch, n=self._nqubits, format=format, jittable=jittable)
+        return sample2all(
+            sample=ch, n=self._nqubits, format=format, jittable=jittable, d=self._d
+        )
 
     def sample_expectation_ps(
         self,
@@ -878,9 +955,9 @@ def replace_inputs(self, inputs: Tensor) -> None:
         """
         inputs = backend.reshape(inputs, [-1])
         N = inputs.shape[0]
-        n = int(np.log(N) / np.log(2))
+        n = _infer_num_sites(N, self._d)
         assert n == self._nqubits
-        inputs = backend.reshape(inputs, [2 for _ in range(n)])
+        inputs = backend.reshape(inputs, [self._d for _ in range(n)])
         if self.inputs is not None:
             self._nodes[0].tensor = inputs
             if self.is_dm:
@@ -911,9 +988,9 @@ def cond_measurement(self, index: int, status: Optional[float] = None) -> Tensor
 
 
 
-        :param index: the qubit for the z-basis measurement
+        :param index: the site index for the Z-basis measurement
         :type index: int
-        :return: 0 or 1 for z measurement on up and down freedom
+        :return: 0 or 1 for Z-basis measurement outcome
         :rtype: Tensor
         """
         return self.general_kraus(  # type: ignore
@@ -992,8 +1069,8 @@ def get_quvector(self) -> QuVector:
 
     def projected_subsystem(self, traceout: Tensor, left: Tuple[int, ...]) -> Tensor:
         """
-        remaining wavefunction or density matrix on qubits in left, with other qubits
-        fixed in 0 or 1 indicated by traceout
+        remaining wavefunction or density matrix on sites in ``left``, with other sites
+        fixed to given digits (0..d-1) as indicated by ``traceout``
 
         :param traceout: can be jitted
         :type traceout: Tensor
@@ -1002,15 +1079,19 @@ def projected_subsystem(self, traceout: Tensor, left: Tuple[int, ...]) -> Tensor
         :return: _description_
         :rtype: Tensor
         """
-        end0, end1 = gates.array_to_tensor(np.array([1.0, 0]), np.array([0, 1.0]))
+
+        def _basis_gate(k_tensor: Any) -> Gate:
+            vec = backend.one_hot(backend.cast(k_tensor, "int32"), self._d)
+            vec = backend.cast(vec, dtypestr)
+            return Gate(vec)
+
         traceout = backend.cast(traceout, dtypestr)
         nodes, front = self._copy()
         L = self._nqubits
         edges = []
         for i in range(len(traceout)):
             if i not in left:
-                b = traceout[i]
-                n = gates.Gate((1 - b) * end0 + b * end1)
+                n = _basis_gate(traceout[i])
                 nodes.append(n)
                 front[i] ^ n[0]
             else:
@@ -1019,8 +1100,7 @@ def projected_subsystem(self, traceout: Tensor, left: Tuple[int, ...]) -> Tensor
         if self.is_dm:
             for i in range(len(traceout)):
                 if i not in left:
-                    b = traceout[i]
-                    n = gates.Gate((1 - b) * end0 + b * end1)
+                    n = _basis_gate(traceout[i])
                     nodes.append(n)
                     front[i + L] ^ n[0]
                 else:

From 209153070d34ce6c69bceec5fa8cb2ec83931dc5 Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Wed, 27 Aug 2025 13:35:05 +0800
Subject: [PATCH 12/55] Adjust system to qudit systems.

---
 tensorcircuit/densitymatrix.py | 29 +++++++++++++++++++++--------
 1 file changed, 21 insertions(+), 8 deletions(-)

diff --git a/tensorcircuit/densitymatrix.py b/tensorcircuit/densitymatrix.py
index f819e058..0a30b707 100644
--- a/tensorcircuit/densitymatrix.py
+++ b/tensorcircuit/densitymatrix.py
@@ -17,7 +17,7 @@
 from .circuit import Circuit
 from .cons import backend, contractor, dtypestr
 from .basecircuit import BaseCircuit
-from .quantum import QuOperator
+from .quantum import QuOperator, _infer_num_sites
 
 Gate = gates.Gate
 Tensor = Any
@@ -29,12 +29,14 @@ class DMCircuit(BaseCircuit):
     def __init__(
         self,
         nqubits: int,
+        dim: Optional[int] = None,
         empty: bool = False,
         inputs: Optional[Tensor] = None,
         mps_inputs: Optional[QuOperator] = None,
         dminputs: Optional[Tensor] = None,
         mpo_dminputs: Optional[QuOperator] = None,
         split: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
     ) -> None:
         """
         The density matrix simulator based on tensornetwork engine.
@@ -55,6 +57,13 @@ def __init__(
             ``max_singular_values`` and ``max_truncation_err``.
         :type split: Optional[Dict[str, Any]]
         """
+        self._d = 2 if dim is None else dim
+        if not kwargs.get("qudit", False) and self._d != 2:
+            raise ValueError(
+                f"Circuit only supports qubits (dim=2). "
+                f"You passed dim={self._d}. Please use `QuditCircuit` instead."
+            )
+
         if not empty:
             if (
                 (inputs is None)
@@ -73,9 +82,9 @@ def __init__(
                 inputs = backend.cast(inputs, dtype=dtypestr)
                 inputs = backend.reshape(inputs, [-1])
                 N = inputs.shape[0]
-                n = int(np.log(N) / np.log(2))
+                n = _infer_num_sites(N, self._d)
                 assert n == nqubits
-                inputs = backend.reshape(inputs, [2 for _ in range(n)])
+                inputs = backend.reshape(inputs, [self._d for _ in range(n)])
                 inputs_gate = Gate(inputs)
                 self._nodes = [inputs_gate]
                 self.coloring_nodes(self._nodes)
@@ -94,7 +103,9 @@ def __init__(
             elif dminputs is not None:
                 dminputs = backend.convert_to_tensor(dminputs)
                 dminputs = backend.cast(dminputs, dtype=dtypestr)
-                dminputs = backend.reshape(dminputs, [2 for _ in range(2 * nqubits)])
+                dminputs = backend.reshape(
+                    dminputs, [self._d for _ in range(2 * nqubits)]
+                )
                 dminputs_gate = Gate(dminputs)
                 nodes = [dminputs_gate]
                 self._front = [dminputs_gate.get_edge(i) for i in range(2 * nqubits)]
@@ -217,7 +228,7 @@ def apply_general_kraus(
             dd = dmc.densitymatrix()
             circuits.append(dd)
         tensor = reduce(add, circuits)
-        tensor = backend.reshape(tensor, [2 for _ in range(2 * self._nqubits)])
+        tensor = backend.reshape(tensor, [self._d for _ in range(2 * self._nqubits)])
         self._nodes = [Gate(tensor)]
         dangling = [e for e in self._nodes[0]]
         self._front = dangling
@@ -255,7 +266,9 @@ def densitymatrix(self, check: bool = False, reuse: bool = True) -> Tensor:
             t = contractor(nodes, output_edge_order=d_edges)
         else:
             t = nodes[0]
-        dm = backend.reshape(t.tensor, shape=[2**self._nqubits, 2**self._nqubits])
+        dm = backend.reshape(
+            t.tensor, shape=[self._d**self._nqubits, self._d**self._nqubits]
+        )
         if check:
             self.check_density_matrix(dm)
         return dm
@@ -274,7 +287,7 @@ def wavefunction(self) -> Tensor:
         dm = self.densitymatrix()
         e, v = backend.eigh(dm)
         np.testing.assert_allclose(
-            e[:-1], backend.zeros([2**self._nqubits - 1]), atol=1e-5
+            e[:-1], backend.zeros([self._d**self._nqubits - 1]), atol=1e-5
         )
         return v[:, -1]
 
@@ -375,7 +388,7 @@ def apply_general_kraus(
         #     index = [index[0] for _ in range(len(kraus))]
         super_op = kraus_to_super_gate(kraus)
         nlegs = 4 * len(index)
-        super_op = backend.reshape(super_op, [2 for _ in range(nlegs)])
+        super_op = backend.reshape(super_op, [self._d for _ in range(nlegs)])
         super_op = Gate(super_op)
         o2i = int(nlegs / 2)
         r2l = int(nlegs / 4)

From db7b5487970b4f1f09312be7266fd9cc02df6cd2 Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Wed, 27 Aug 2025 13:55:00 +0800
Subject: [PATCH 13/55] Adjust system to qudit systems. Passes all tests in
 test_mpscircuit.py

---
 tensorcircuit/mpscircuit.py | 166 +++++++++++++++++++++++-------------
 1 file changed, 107 insertions(+), 59 deletions(-)

diff --git a/tensorcircuit/mpscircuit.py b/tensorcircuit/mpscircuit.py
index dcd8f0c8..d2e9be87 100644
--- a/tensorcircuit/mpscircuit.py
+++ b/tensorcircuit/mpscircuit.py
@@ -8,6 +8,7 @@
 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
@@ -17,6 +18,7 @@
 from .quantum import QuOperator, QuVector, extract_tensors_from_qop
 from .mps_base import FiniteMPS
 from .abstractcircuit import AbstractCircuit
+from .basecircuit import _decode_basis_label
 from .utils import arg_alias
 
 Gate = gates.Gate
@@ -87,16 +89,20 @@ class MPSCircuit(AbstractCircuit):
     def __init__(
         self,
         nqubits: int,
+        dim: Optional[int] = None,
         center_position: Optional[int] = None,
         tensors: Optional[Sequence[Tensor]] = None,
         wavefunction: Optional[Union[QuVector, Tensor]] = None,
         split: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
     ) -> None:
         """
         MPSCircuit object based on state simulator.
 
         :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``
@@ -109,6 +115,13 @@ def __init__(
         :param split: Split rules
         :type split: Any
         """
+        self._d = 2 if dim is None else dim
+        if not kwargs.get("qudit", False) and self._d != 2:
+            raise ValueError(
+                f"Circuit only supports qubits (dim=2). "
+                f"You passed dim={self._d}. Please use `QuditCircuit` instead."
+            )
+
         self.circuit_param = {
             "nqubits": nqubits,
             "center_position": center_position,
@@ -137,7 +150,9 @@ def __init__(
                         wavefunction, split=self.split
                     )
             else:  # full wavefunction
-                tensors = self.wavefunction_to_tensors(wavefunction, split=self.split)
+                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
@@ -151,8 +166,13 @@ def __init__(
                 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:
@@ -370,42 +390,57 @@ def gate_to_MPO(
         #      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 AssertionError("`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)
-        # 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)
-        # reorder the gate according to the site positions
-        gate = backend.transpose(gate, tuple(argsort))
-        index = index[argsort]  # type: ignore
-        # split the gate into tensors assuming they are adjacent
-        main_tensors = cls.wavefunction_to_tensors(gate, dim_phys=4, 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
+        in_dims = tuple(backend.shape_tuple(gate))[:nindex]
+        d = int(in_dims[0])
+        dim_phys_mpo = d * d
+
+        order = tuple(
+            np.arange(2 * nindex).reshape(2, nindex).T.flatten().tolist()
+        )
+        gate = backend.transpose(gate, order)
+
+        index_arr = np.array(index, dtype=int) - index_left
+        pair_order = np.argsort(index_arr)
+
+        pair_axis_perm = np.ravel(
+            np.column_stack([2 * pair_order, 2 * pair_order + 1])
+        ).astype(int)
+        pair_axis_perm = tuple(pair_axis_perm.tolist())  # type: ignore
+        gate = backend.transpose(gate, pair_axis_perm)
+        index_arr = index_arr[pair_order]  # type: ignore
+
+        gate = backend.reshape(gate, (dim_phys_mpo,) * nindex)
+        main_tensors = cls.wavefunction_to_tensors(
+            gate, dim_phys=dim_phys_mpo, norm=False
+        )
+
+        tensors: list[Tensor] = []
+        previous_i: Optional[int] = None
+
+        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)
-                    )
-                    tensors.append(backend.convert_to_tensor(I))
-            nleft, _, nright = main_tensor.shape
-            tensor = backend.reshape(main_tensor, (nleft, 2, 2, nright))
+                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 * d, dtype=backend.dtype(tensors[-1]))
+                    I4 = backend.reshape(eye2d, (bond_dim, d, bond_dim, d))
+                    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), d, d, int(nright)))
             tensors.append(tensor)
-            previous_i = i
+            previous_i = int(i)
+
         return tensors, index_left
 
     @classmethod
@@ -448,15 +483,15 @@ def reduce_tensor_dimension(
         """
         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(
@@ -550,10 +585,11 @@ def apply_MPO(
         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
@@ -660,8 +696,7 @@ def mid_measurement(self, index: int, keep: int = 0) -> None:
         :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]]),
@@ -692,7 +727,7 @@ def is_valid(self) -> bool:
     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]:
@@ -710,6 +745,7 @@ def wavefunction_to_tensors(
         :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))
@@ -768,10 +804,16 @@ def copy_without_tensor(self) -> "MPSCircuit":
         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
 
@@ -815,7 +857,8 @@ def normalize(self) -> None:
 
     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, self._d, self._nqubits)
+        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:
@@ -873,6 +916,7 @@ def slice(self, begin: int, end: int) -> "MPSCircuit":
 
         mps = self.__class__(
             nqubits,
+            dim=self._d,
             tensors=tensors,
             center_position=center_position,
             split=self.split.copy(),
@@ -1000,8 +1044,6 @@ def measure(
         # 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)
@@ -1015,20 +1057,26 @@ def measure(
                 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]
             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
+
+            cdf = backend.cumsum(ps)
+            choice = backend.sum(backend.cast(r >= cdf, "int32"))
+
+            choice_f = backend.cast(choice, dtypestr)
+            sample.append(choice_f)
+
+            m = backend.zeros((ps.shape[0],), dtype=dtypestr)
+            m = backend.scatter(
+                m,
+                backend.convert_to_tensor([[backend.cast(choice, "int32")]]),
+                backend.convert_to_tensor(np.array([1.0], dtype=dtypestr)),
+            )
+
+            p = p * backend.sum(ps * backend.cast(m, dtype=rdtypestr))
             mps._mps.tensors[site] = backend.einsum("iaj,a->ij", tensor, m)[:, None, :]
         sample = backend.stack(sample)
         sample = backend.real(sample)

From c18944f06a56b20748bda827890e925900c3349f Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Wed, 27 Aug 2025 15:10:26 +0800
Subject: [PATCH 14/55] black .

---
 tensorcircuit/mpscircuit.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/tensorcircuit/mpscircuit.py b/tensorcircuit/mpscircuit.py
index d2e9be87..8cd953a2 100644
--- a/tensorcircuit/mpscircuit.py
+++ b/tensorcircuit/mpscircuit.py
@@ -404,9 +404,7 @@ def gate_to_MPO(
         d = int(in_dims[0])
         dim_phys_mpo = d * d
 
-        order = tuple(
-            np.arange(2 * nindex).reshape(2, nindex).T.flatten().tolist()
-        )
+        order = tuple(np.arange(2 * nindex).reshape(2, nindex).T.flatten().tolist())
         gate = backend.transpose(gate, order)
 
         index_arr = np.array(index, dtype=int) - index_left

From c512a59ad803734a40b7181e8c45322b657a720c Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Thu, 28 Aug 2025 15:49:22 +0800
Subject: [PATCH 15/55] Considering the existence of resetting the circuit in
 the line (such as matrix(), get_quoperator() method), add a _qudit attribute
 to abstractcircuit.py, defaulting to False.

---
 tensorcircuit/abstractcircuit.py | 1 +
 tensorcircuit/circuit.py         | 7 ++++---
 tensorcircuit/densitymatrix.py   | 3 ++-
 tensorcircuit/mpscircuit.py      | 4 ++--
 4 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/tensorcircuit/abstractcircuit.py b/tensorcircuit/abstractcircuit.py
index 6dd6d07d..1e8594c6 100644
--- a/tensorcircuit/abstractcircuit.py
+++ b/tensorcircuit/abstractcircuit.py
@@ -69,6 +69,7 @@
 class AbstractCircuit:
     _nqubits: int
     _d: int = 2
+    _qudit: bool = False
     _qir: List[Dict[str, Any]]
     _extra_qir: List[Dict[str, Any]]
     inputs: Tensor
diff --git a/tensorcircuit/circuit.py b/tensorcircuit/circuit.py
index 01a53e08..fd9c5e92 100644
--- a/tensorcircuit/circuit.py
+++ b/tensorcircuit/circuit.py
@@ -67,7 +67,8 @@ def __init__(
         :type split: Optional[Dict[str, Any]]
         """
         self._d = 2 if dim is None else dim
-        if not kwargs.get("qudit", False) and self._d != 2:
+        self._qudit: bool = kwargs.get("qudit", False)
+        if not self._qudit and self._d != 2:
             raise ValueError(
                 f"Circuit only supports qubits (dim=2). "
                 f"You passed dim={self._d}. Please use `QuditCircuit` instead."
@@ -739,7 +740,7 @@ def get_quoperator(self) -> QuOperator:
         :rtype: QuOperator
         """
         mps = identity([self._d for _ in range(self._nqubits)])
-        c = Circuit(self._nqubits, self._d)
+        c = Circuit(self._nqubits, self._d, qudit=self._qudit)
         ns, es = self._copy()
         c._nodes = ns
         c._front = es
@@ -760,7 +761,7 @@ def matrix(self) -> Tensor:
         :rtype: Tensor
         """
         mps = identity([self._d for _ in range(self._nqubits)])
-        c = Circuit(self._nqubits, self._d)
+        c = Circuit(self._nqubits, self._d, qudit=self._qudit)
         ns, es = self._copy()
         c._nodes = ns
         c._front = es
diff --git a/tensorcircuit/densitymatrix.py b/tensorcircuit/densitymatrix.py
index 0a30b707..514b5264 100644
--- a/tensorcircuit/densitymatrix.py
+++ b/tensorcircuit/densitymatrix.py
@@ -58,7 +58,8 @@ def __init__(
         :type split: Optional[Dict[str, Any]]
         """
         self._d = 2 if dim is None else dim
-        if not kwargs.get("qudit", False) and self._d != 2:
+        self._qudit: bool = kwargs.get("qudit", False)
+        if not self._qudit and self._d != 2:
             raise ValueError(
                 f"Circuit only supports qubits (dim=2). "
                 f"You passed dim={self._d}. Please use `QuditCircuit` instead."
diff --git a/tensorcircuit/mpscircuit.py b/tensorcircuit/mpscircuit.py
index 8cd953a2..b8c511ef 100644
--- a/tensorcircuit/mpscircuit.py
+++ b/tensorcircuit/mpscircuit.py
@@ -115,8 +115,8 @@ def __init__(
         :param split: Split rules
         :type split: Any
         """
-        self._d = 2 if dim is None else dim
-        if not kwargs.get("qudit", False) and self._d != 2:
+        self._qudit: bool = kwargs.get("qudit", False)
+        if not self._qudit and self._d != 2:
             raise ValueError(
                 f"Circuit only supports qubits (dim=2). "
                 f"You passed dim={self._d}. Please use `QuditCircuit` instead."

From 0f319e9e7cefcbfbda029980f67e4fb44ce3c10d Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Thu, 28 Aug 2025 16:39:02 +0800
Subject: [PATCH 16/55] Fixed a potential error when the backend is TensorFlow.

---
 tensorcircuit/basecircuit.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/tensorcircuit/basecircuit.py b/tensorcircuit/basecircuit.py
index c84acf6c..14ca9156 100644
--- a/tensorcircuit/basecircuit.py
+++ b/tensorcircuit/basecircuit.py
@@ -482,10 +482,9 @@ def measure_jit(
                 tiny_r = backend.cast(backend.convert_to_tensor(1e-12), rdtypestr)
                 pu = backend.real(backend.diagonal(rho))
                 pu = backend.clip(pu, zero_r, one_r)
-                d = backend.shape_tuple(pu)[-1]
                 pu = pu + tiny_r * (
-                    backend.ones((d,), dtype=rdtypestr)
-                    / backend.cast(backend.convert_to_tensor(float(d)), rdtypestr)
+                    backend.ones((self._d,), dtype=rdtypestr)
+                    / backend.cast(backend.convert_to_tensor(float(self._d)), rdtypestr)
                 )
                 pu = pu / backend.sum(pu)
                 cdf = backend.cumsum(pu)
@@ -502,11 +501,11 @@ def measure_jit(
                     r = backend.clip(r, zero_r, one_r - tiny_r)
                 else:
                     r = backend.real(backend.cast(status[k], rdtypestr))
-                k_out = backend.searchsorted(cdf, r, side="right")
+                k_out = backend.sum(backend.cast(cdf <= r, "int32"))
                 k_out = backend.clip(
                     k_out,
                     backend.cast(backend.convert_to_tensor(0), "int32"),
-                    backend.cast(backend.convert_to_tensor(d - 1), "int32"),
+                    backend.cast(backend.convert_to_tensor(self._d - 1), "int32"),
                 )
                 sample.append(backend.cast(k_out, rdtypestr))
                 p = p * backend.cast(pu[k_out], rdtypestr)

From 2a3e1aa7d6f6971ee1cfa2c7702e4e5453ac28cf Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Fri, 29 Aug 2025 15:32:14 +0800
Subject: [PATCH 17/55] removed unnecessary codes.

---
 tensorcircuit/basecircuit.py | 8 --------
 1 file changed, 8 deletions(-)

diff --git a/tensorcircuit/basecircuit.py b/tensorcircuit/basecircuit.py
index 14ca9156..2318f080 100644
--- a/tensorcircuit/basecircuit.py
+++ b/tensorcircuit/basecircuit.py
@@ -491,14 +491,6 @@ def measure_jit(
                 if status is None:
                     r = backend.implicit_randu()[0]
                     r = backend.real(backend.cast(r, rdtypestr))
-                    phi = backend.cast(
-                        backend.convert_to_tensor(0.6180339887498948), rdtypestr
-                    )
-                    r = r + phi * backend.cast(
-                        backend.convert_to_tensor(k + 1), rdtypestr
-                    )
-                    r = r - backend.floor(r)
-                    r = backend.clip(r, zero_r, one_r - tiny_r)
                 else:
                     r = backend.real(backend.cast(status[k], rdtypestr))
                 k_out = backend.sum(backend.cast(cdf <= r, "int32"))

From 91aea863e3814cd94cfb010730489a3d8f084328 Mon Sep 17 00:00:00 2001
From: refraction-ray <kcanamgal@foxmail.com>
Date: Sat, 30 Aug 2025 17:23:22 +0800
Subject: [PATCH 18/55] fix #41

# Conflicts:
#	tensorcircuit/basecircuit.py
#	tensorcircuit/quantum.py
---
 CHANGELOG.md             |  2 ++
 tensorcircuit/quantum.py | 11 +++++++++++
 2 files changed, 13 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2d1ba410..0be7e346 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -40,6 +40,8 @@
 
 - Fix to use `status` for `circuit.sample` when `allow_state=True`.
 
+- Fix sample bug when number of qubit exceeding 32.
+
 ### Changed
 
 - The order of arguments of `tc.timeevol.ed_evol` are changed for consistent interface with other evolution methods.
diff --git a/tensorcircuit/quantum.py b/tensorcircuit/quantum.py
index 3454e27a..891a9c23 100644
--- a/tensorcircuit/quantum.py
+++ b/tensorcircuit/quantum.py
@@ -9,6 +9,7 @@
 import os
 from functools import partial, reduce
 from operator import matmul, mul, or_
+from collections import Counter
 from typing import (
     Any,
     Callable,
@@ -2971,6 +2972,16 @@ def sample2all(
     :rtype: Any
     """
     d = 2 if d is None else int(d)
+    if n > 32:
+        assert (
+            len(backend.shape_tuple(sample)) == 2
+        ), "n>32 is only supported for ``sample_bin``"
+        if format == "sample_bin":
+            return sample
+        if format == "count_dict_bin":
+            binary_strings = ["".join(map(str, shot)) for shot in sample]
+            return dict(Counter(binary_strings))
+        raise ValueError(f"n={n} is too large for measurement representaion: {format}")
 
     if len(backend.shape_tuple(sample)) == 1:
         sample_int = sample

From 48fb567e0ad6780db3b6dae5cff201242c0cca98 Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Sun, 31 Aug 2025 15:51:33 +0800
Subject: [PATCH 19/55] fix bug

---
 tensorcircuit/quantum.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorcircuit/quantum.py b/tensorcircuit/quantum.py
index 891a9c23..b36ce7c1 100644
--- a/tensorcircuit/quantum.py
+++ b/tensorcircuit/quantum.py
@@ -2485,7 +2485,7 @@ def entanglement_negativity(
     :return: _description_
     :rtype: Tensor
     """
-    rhot = partial_transpose(rho, transposed_sites)
+    rhot = partial_transpose(rho, transposed_sites, d=d)
     es = backend.eigvalsh(rhot)
     rhot_m = backend.sum(backend.abs(es))
     return (rhot_m - 1.0) / 2.0

From d9cbed2aa728198f1fddde0b9dfedccf3cad3218 Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Sun, 31 Aug 2025 17:17:01 +0800
Subject: [PATCH 20/55] add floor_divide() to all the backends.

---
 tensorcircuit/backends/abstract_backend.py   | 38 +++++++++++++++++---
 tensorcircuit/backends/jax_backend.py        |  3 ++
 tensorcircuit/backends/numpy_backend.py      |  3 ++
 tensorcircuit/backends/pytorch_backend.py    |  3 ++
 tensorcircuit/backends/tensorflow_backend.py |  3 ++
 5 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/tensorcircuit/backends/abstract_backend.py b/tensorcircuit/backends/abstract_backend.py
index e7de5654..ac3f3d9f 100644
--- a/tensorcircuit/backends/abstract_backend.py
+++ b/tensorcircuit/backends/abstract_backend.py
@@ -865,7 +865,35 @@ def mod(self: Any, x: Tensor, y: Tensor) -> Tensor:
             "Backend '{}' has not implemented `mod`.".format(self.name)
         )
 
-    def floor(self: Any, x: Tensor) -> Tensor:
+    def floor_divide(self: Any, x: Tensor, y: Tensor) -> Tensor:
+        r"""
+        Compute the element-wise floor division of two tensors.
+
+        This operation returns a new tensor containing the result of
+        dividing `x` by `y` and rounding each element down towards
+        negative infinity. The semantics are equivalent to the Python
+        `//` operator:
+
+            result[i] = floor(x[i] / y[i])
+
+        Broadcasting is supported according to the backend's rules.
+
+        :param x: Dividend tensor.
+        :type x: Tensor
+        :param y: Divisor tensor, must be broadcastable with `x`.
+        :type y: Tensor
+        :return: A tensor with the broadcasted shape of `x` and `y`,
+                 where each element is the floored result of the division.
+        :rtype: Tensor
+
+        :raises NotImplementedError: If the backend does not provide an
+                                     implementation for `floor_divide`.
+        """
+        raise NotImplementedError(
+            "Backend '{}' has not implemented `floor_divide`.".format(self.name)
+        )
+
+    def floor(self: Any, a: Tensor) -> Tensor:
         """
         Compute the element-wise floor of the input tensor.
 
@@ -873,10 +901,10 @@ def floor(self: Any, x: Tensor) -> Tensor:
         less than or equal to each element of the input tensor,
         i.e. it rounds each value down towards negative infinity.
 
-        :param x: Input tensor containing numeric values.
-        :type x: Tensor
-        :return: A tensor with the same shape as `x`, where each element
-                 is the floored value of the corresponding element in `x`.
+        :param a: Input tensor containing numeric values.
+        :type a: Tensor
+        :return: A tensor with the same shape as `a`, where each element
+                 is the floored value of the corresponding element in `a`.
         :rtype: Tensor
 
         :raises NotImplementedError: If the backend does not provide an
diff --git a/tensorcircuit/backends/jax_backend.py b/tensorcircuit/backends/jax_backend.py
index 880ad4e8..4ff5cb33 100644
--- a/tensorcircuit/backends/jax_backend.py
+++ b/tensorcircuit/backends/jax_backend.py
@@ -352,6 +352,9 @@ def mod(self, x: Tensor, y: Tensor) -> Tensor:
     def floor(self, a: Tensor) -> Tensor:
         return jnp.floor(a)
 
+    def floor_divide(self, x: Tensor, y: Tensor) -> Tensor:
+        return jnp.floor_divide(x, y)
+
     def clip(self, a: Tensor, a_min: Tensor, a_max: Tensor) -> Tensor:
         return jnp.clip(a, a_min, a_max)
 
diff --git a/tensorcircuit/backends/numpy_backend.py b/tensorcircuit/backends/numpy_backend.py
index 02669ca7..ce165b7d 100644
--- a/tensorcircuit/backends/numpy_backend.py
+++ b/tensorcircuit/backends/numpy_backend.py
@@ -250,6 +250,9 @@ def arange(self, start: int, stop: Optional[int] = None, step: int = 1) -> Tenso
     def mod(self, x: Tensor, y: Tensor) -> Tensor:
         return np.mod(x, y)
 
+    def floor_divide(self, x: Tensor, y: Tensor) -> Tensor:
+        return np.floor_divide(x, y)
+
     def floor(self, a: Tensor) -> Tensor:
         return np.floor(a)
 
diff --git a/tensorcircuit/backends/pytorch_backend.py b/tensorcircuit/backends/pytorch_backend.py
index dba8e1a0..a54ea206 100644
--- a/tensorcircuit/backends/pytorch_backend.py
+++ b/tensorcircuit/backends/pytorch_backend.py
@@ -429,6 +429,9 @@ def arange(self, start: int, stop: Optional[int] = None, step: int = 1) -> Tenso
     def mod(self, x: Tensor, y: Tensor) -> Tensor:
         return torchlib.fmod(x, y)
 
+    def floor_divide(self, x: Tensor, y: Tensor) -> Tensor:
+        return torchlib.floor_divide(x, y)
+
     def floor(self, a: Tensor) -> Tensor:
         return torchlib.floor(a)
 
diff --git a/tensorcircuit/backends/tensorflow_backend.py b/tensorcircuit/backends/tensorflow_backend.py
index 82a3e3d6..4818d2dc 100644
--- a/tensorcircuit/backends/tensorflow_backend.py
+++ b/tensorcircuit/backends/tensorflow_backend.py
@@ -581,6 +581,9 @@ def floor(self, a: Tensor) -> Tensor:
             return a
         return tf.math.floor(a)
 
+    def floor_divide(self, x: Tensor, y: Tensor) -> Tensor:
+        return tf.math.floordiv(x, y)
+
     def concat(self, a: Sequence[Tensor], axis: int = 0) -> Tensor:
         return tf.concat(a, axis=axis)
 

From 3d05a5fbb4689de24f0524010be849b94efa6e7a Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Sun, 31 Aug 2025 17:17:28 +0800
Subject: [PATCH 21/55] add test for floor_divide().

---
 tests/test_backends.py | 56 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 56 insertions(+)

diff --git a/tests/test_backends.py b/tests/test_backends.py
index 01422429..265a15f2 100644
--- a/tests/test_backends.py
+++ b/tests/test_backends.py
@@ -1335,3 +1335,59 @@ def test_backend_where(backend):
     result = tc.backend.where(condition, x, y)
     expected = tc.backend.convert_to_tensor([1, 5, 3])
     np.testing.assert_allclose(result, expected)
+
+
+@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb"), lf("torchb")])
+def test_floor_divide_various_cases(backend):
+    r"""
+    Single test covering:
+      - basic positive integers
+      - negative dividends/divisors
+      - broadcasting
+      - floating inputs (matches floor semantics)
+    Ensures both operands are converted to the active backend's native tensor/array
+    to avoid type errors in Torch/JAX.
+    """
+
+    def to_backend(z):
+        if hasattr(tc.backend, "asarray"):
+            return tc.backend.asarray(z)
+
+        name = getattr(tc.backend, "name", "").lower()
+        try:
+            if "torch" in name:
+                import torch
+                return torch.as_tensor(z)
+            if "jax" in name:
+                import jax.numpy as jnp
+                return jnp.asarray(z)
+            if "tf" in name or "tensorflow" in name:
+                import tensorflow as tf
+                return tf.convert_to_tensor(z)
+        except Exception:
+            pass
+        return np.asarray(z)
+
+    out = tc.backend.floor_divide(to_backend([7, 8, 9]), to_backend(2))
+    np.testing.assert_array_equal(np.array(out), np.array([3, 4, 4]))
+
+    out = tc.backend.floor_divide(to_backend([-3, -4]), to_backend(2))
+    np.testing.assert_array_equal(np.array(out), np.array([-2, -2]))
+
+    out = tc.backend.floor_divide(to_backend([3, 4]), to_backend(-2))
+    np.testing.assert_array_equal(np.array(out), np.array([-2, -2]))
+
+    out = tc.backend.floor_divide(to_backend([-3, -4]), to_backend(-2))
+    np.testing.assert_array_equal(np.array(out), np.array([1, 2]))
+
+    x = to_backend([[10, 20], [30, 40]])
+    y = to_backend([3, 5])
+    expected = np.array([[10, 20], [30, 40]]) // np.array([3, 5])
+    out = tc.backend.floor_divide(x, y)
+    np.testing.assert_array_equal(np.array(out), expected)
+
+    xf = to_backend([7.9, 8.1, -3.5])
+    yf = to_backend([2.0, 2.0, 2.0])
+    expectedf = np.floor_divide(np.array([7.9, 8.1, -3.5]), np.array([2.0, 2.0, 2.0]))
+    out = tc.backend.floor_divide(xf, yf)
+    np.testing.assert_array_equal(np.array(out), expectedf)
\ No newline at end of file

From c7227e265383beb4f0d116a9fa5adb3d2a35791d Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Sun, 31 Aug 2025 17:18:50 +0800
Subject: [PATCH 22/55] remove redundant codes

---
 tensorcircuit/results/counts.py | 67 +++++++--------------------------
 1 file changed, 13 insertions(+), 54 deletions(-)

diff --git a/tensorcircuit/results/counts.py b/tensorcircuit/results/counts.py
index b38dada8..a921f209 100644
--- a/tensorcircuit/results/counts.py
+++ b/tensorcircuit/results/counts.py
@@ -90,7 +90,7 @@ def marginal_count(count: ct, keep_list: Sequence[int]) -> ct:
     return reverse_count(ncount)
 
 
-def count2vec(count: ct, normalization: bool = True) -> Tensor:
+def count2vec(count: ct, normalization: bool = True, dim: Optional[int] = 2) -> Tensor:
     """
     Convert a dictionary of counts (with string keys) to a probability/count vector.
 
@@ -104,6 +104,8 @@ def count2vec(count: ct, normalization: bool = True) -> Tensor:
     :type count: ct
     :param normalization: Whether to normalize the counts to probabilities, defaults to True
     :type normalization: bool, optional
+    :param dim: Dimensionality of the vector, defaults to 2
+    :type dim: int, optional
     :return: Probability vector as numpy array
     :rtype: Tensor
 
@@ -112,37 +114,20 @@ def count2vec(count: ct, normalization: bool = True) -> Tensor:
     >>> count2vec({"00": 2, "10": 3, "11": 5})
     array([0.2, 0. , 0.3, 0.5])
     """
-    if not count:
-        return np.array([], dtype=float)
-
-    sample_key = next(iter(count)).upper()
-    n = len(sample_key)
-    d = 0
-    for k in count:
-        s = k.upper()
-        if len(s) != n:
-            raise ValueError(
-                f"The length of all keys should be the same ({n}), received '{k}'."
-            )
-        for ch in s:
-            if ch not in _ALPHABET:
-                raise ValueError(
-                    f"Key '{k}' contains illegal character '{ch}' (only 0-9A-Z are allowed)."
-                )
-            d = max(d, _ALPHABET.index(ch) + 1)
-    if d < 2:
-        raise ValueError(f"Inferred local dimension d={d} is illegal (must be >=2).")
 
     def parse_key(_k: str) -> List[int]:
         return [_ALPHABET.index(_ch) for _ch in _k.upper()]
 
-    size = d**n
-    prob = np.zeros(size, dtype=float)
+    if not count:
+        return np.array([], dtype=float)
+
+    n = len(next(iter(count)).upper())
+    prob = np.zeros(dim**n, dtype=float)
     shots = float(sum(count.values())) if normalization else 1.0
     if shots == 0:
         return prob
 
-    powers = [d**p for p in range(n)][::-1]
+    powers = [dim**p for p in range(n)][::-1]
     for k, v in count.items():
         digits = parse_key(k)
         idx = sum(dig * p for dig, p in zip(digits, powers))
@@ -151,48 +136,22 @@ def parse_key(_k: str) -> List[int]:
     return prob
 
 
-def vec2count(vec: Tensor, prune: bool = False) -> ct:
+def vec2count(vec: Tensor, prune: bool = False, dim: Optional[int] = 2) -> ct:
     """
     Map a count/probability vector of length D to a dictionary with base-d string keys (0-9A-Z).
     Only generate string keys when d ≤ 36; if d is inferred to be > 36, raise a NotImplementedError.
 
     :param vec: A one-dimensional vector of length D = d**n
     :param prune: Whether to prune near-zero elements (threshold 1e-8)
+    :param dim: Dimensionality of the vector, defaults to 2
     :return: {base-d string key: value}, key length n
     """
     from ..quantum import count_vector2dict, _infer_num_sites
 
     if isinstance(vec, list):
         vec = np.array(vec)
-    vec = np.asarray(vec)
-    if vec.ndim != 1:
-        raise ValueError("vec2count expects a one-dimensional vector.")
-
-    D = int(vec.shape[0])
-    if D <= 0:
-        return {}
-
-    def _is_power_of_two(x: int) -> bool:
-        return x > 0 and (x & (x - 1)) == 0
-
-    if _is_power_of_two(D):
-        n = int(np.log(D) / np.log(2) + 1e-9)
-        d: Optional[int] = 2
-    else:
-        d = n = None
-        upper = int(np.sqrt(D)) + 1
-        for d_try in range(2, max(upper, 3)):
-            try:
-                n_try = _infer_num_sites(D, d_try)
-            except ValueError:
-                continue
-            d, n = d_try, n_try
-            break
-        if d is None:
-            d, n = D, 1
-
-    c: ct = count_vector2dict(vec, n, key="bin", d=d)  # type: ignore
-
+    n = int(np.log(int(vec.shape[0])) / np.log(dim) + 1e-9)
+    c: ct = count_vector2dict(vec, n, key="bin", d=dim)  # type: ignore
     if prune:
         c = {k: v for k, v in c.items() if np.abs(v) >= 1e-8}
 

From 9e8876589703ec4e2f139728d772cb95d05f7ed5 Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Sun, 31 Aug 2025 17:19:14 +0800
Subject: [PATCH 23/55] make codes consistent

---
 tensorcircuit/basecircuit.py | 35 ++++++++++++++---------------------
 1 file changed, 14 insertions(+), 21 deletions(-)

diff --git a/tensorcircuit/basecircuit.py b/tensorcircuit/basecircuit.py
index 2318f080..3d6c783f 100644
--- a/tensorcircuit/basecircuit.py
+++ b/tensorcircuit/basecircuit.py
@@ -433,29 +433,22 @@ def measure_jit(
                     m = (1 - sample[i]) * gates.array_to_tensor(
                         np.array([1, 0])
                     ) + sample[i] * gates.array_to_tensor(np.array([0, 1]))
-                    g1 = Gate(m)
-                    g1.id = id(g1)
-                    g1.is_dagger = False
-                    g1.flag = "measurement"
-                    newnodes.append(g1)
-                    g1.get_edge(0) ^ edge1[index[i]]
-                    g2 = Gate(m)
-                    g2.id = id(g2)
-                    g2.is_dagger = True
-                    g2.flag = "measurement"
-                    newnodes.append(g2)
-                    g2.get_edge(0) ^ edge2[index[i]]
                 else:
                     vec = backend.one_hot(backend.cast(sample[i], "int32"), self._d)
-                    v = backend.cast(vec, dtypestr)
-                    m = backend.tensordot(v, v, axes=0)
-                    g = Gate(m)
-                    g.id = id(g)
-                    g.is_dagger = False
-                    g.flag = "measurement"
-                    newnodes.append(g)
-                    g.get_edge(0) ^ edge1[index[i]]
-                    g.get_edge(1) ^ edge2[index[i]]
+                    m = backend.cast(vec, dtypestr)
+                g1 = Gate(m)
+                g1.id = id(g1)
+                g1.is_dagger = False
+                g1.flag = "measurement"
+                newnodes.append(g1)
+                g1.get_edge(0) ^ edge1[index[i]]
+                g2 = Gate(m)
+                g2.id = id(g2)
+                g2.is_dagger = True
+                g2.flag = "measurement"
+                newnodes.append(g2)
+                g2.get_edge(0) ^ edge2[index[i]]
+
             rho = (
                 1
                 / backend.cast(p, dtypestr)

From 855b28c68dd954092334eb8701188c6180c577d0 Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Sun, 31 Aug 2025 17:19:39 +0800
Subject: [PATCH 24/55] fixed a potential bug.

---
 tensorcircuit/circuit.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorcircuit/circuit.py b/tensorcircuit/circuit.py
index fd9c5e92..9f3b0ea7 100644
--- a/tensorcircuit/circuit.py
+++ b/tensorcircuit/circuit.py
@@ -820,7 +820,7 @@ def measure_reference(
             )
             probs = backend.real(backend.diagonal(rho))
             probs /= backend.sum(probs)
-            outcome = np.random.choice(self._d, p=probs)
+            outcome = np.random.choice(self._d, p=backend.cast(probs, dtype="float64"))
 
             sample += _ALPHABET[outcome]
             p *= float(probs[outcome])

From 44f6e0c01ed1c11f8690d2d51b17191297fc572d Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Sun, 31 Aug 2025 17:20:03 +0800
Subject: [PATCH 25/55] fixed a potential bug.

---
 tensorcircuit/quantum.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorcircuit/quantum.py b/tensorcircuit/quantum.py
index b36ce7c1..9a98f5b7 100644
--- a/tensorcircuit/quantum.py
+++ b/tensorcircuit/quantum.py
@@ -2725,7 +2725,7 @@ def sample_int2bin(sample: Tensor, n: int, d: Optional[int] = None) -> Tensor:
         pos = backend.reverse(backend.arange(n))
         base = backend.power(d, pos)
         digits = backend.mod(
-            backend.floor(backend.divide(sample[..., None], base)),  # ⌊sample / d**pos⌋
+            backend.floor_divide(sample[..., None], base),  # ⌊sample / d**pos⌋
             d,
         )
         return backend.cast(digits, "int32")

From cb6dcf689486a9d01595de08415b0ce122685a00 Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Sun, 31 Aug 2025 17:22:40 +0800
Subject: [PATCH 26/55] add tests for d>2 case.

---
 tests/test_results.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/tests/test_results.py b/tests/test_results.py
index 9be81889..35e54383 100644
--- a/tests/test_results.py
+++ b/tests/test_results.py
@@ -49,6 +49,12 @@ def test_merge_count():
 
 def test_count2vec():
     assert counts.vec2count(counts.count2vec(d, normalization=False), prune=True) == d
+    assert (
+        counts.vec2count(
+            counts.count2vec(d, normalization=False, dim=36), prune=True, dim=36
+        )
+        == d
+    )
 
 
 def test_kl():

From 8cbf98e409ca7fa909b6e52bbee72efa305f5200 Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Sun, 31 Aug 2025 17:23:00 +0800
Subject: [PATCH 27/55] add test for d>2 cases.

---
 tests/test_results.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/test_results.py b/tests/test_results.py
index 35e54383..3d4f1199 100644
--- a/tests/test_results.py
+++ b/tests/test_results.py
@@ -7,6 +7,7 @@
 from tensorcircuit.results.readout_mitigation import ReadoutMit
 
 d = {"000": 2, "101": 3, "100": 4}
+d_higher = {"A00": 2, "9BC": 3, "XYZ": 4}
 
 
 def test_marginal_count():

From 8b7b227810418affe990629360a31518ebce7aad Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Mon, 1 Sep 2025 12:33:31 +0800
Subject: [PATCH 28/55] use backend

---
 tensorcircuit/circuit.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorcircuit/circuit.py b/tensorcircuit/circuit.py
index 9f3b0ea7..9220929f 100644
--- a/tensorcircuit/circuit.py
+++ b/tensorcircuit/circuit.py
@@ -820,7 +820,7 @@ def measure_reference(
             )
             probs = backend.real(backend.diagonal(rho))
             probs /= backend.sum(probs)
-            outcome = np.random.choice(self._d, p=backend.cast(probs, dtype="float64"))
+            outcome = backend.implicit_randc(self._d, shape=1, p=probs)
 
             sample += _ALPHABET[outcome]
             p *= float(probs[outcome])

From a7e1577c06a292c75b636f30e4b7f699b0c088d7 Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Mon, 1 Sep 2025 12:52:58 +0800
Subject: [PATCH 29/55] d -> dim

---
 tensorcircuit/basecircuit.py    |  24 ++--
 tensorcircuit/circuit.py        |  20 +--
 tensorcircuit/mpscircuit.py     |  10 +-
 tensorcircuit/quantum.py        | 232 ++++++++++++++++----------------
 tensorcircuit/results/counts.py |   2 +-
 5 files changed, 145 insertions(+), 143 deletions(-)

diff --git a/tensorcircuit/basecircuit.py b/tensorcircuit/basecircuit.py
index 3d6c783f..f63731c6 100644
--- a/tensorcircuit/basecircuit.py
+++ b/tensorcircuit/basecircuit.py
@@ -37,10 +37,10 @@
 Tensor = Any
 
 
-def _decode_basis_label(label: str, d: int, n: int) -> List[int]:
-    if d > 36:
+def _decode_basis_label(label: str, dim: int, n: int) -> List[int]:
+    if dim > 36:
         raise NotImplementedError(
-            f"String basis label supports d<=36 (0–9A–Z). Got d={d}. "
+            f"String basis label supports d<=36 (0–9A–Z). Got dim={dim}. "
             "Use an integer array/tensor of length n instead."
         )
     s = label.upper()
@@ -53,9 +53,9 @@ def _decode_basis_label(label: str, d: int, n: int) -> List[int]:
                 f"Invalid character '{ch}' in basis label (allowed 0–9A–Z)."
             )
         v = _ALPHABET.index(ch)
-        if v >= d:
+        if v >= dim:
             raise ValueError(
-                f"Digit '{ch}' (= {v}) out of range for base-d with d={d}."
+                f"Digit '{ch}' (= {v}) out of range for base-d with dim={dim}."
             )
         digits.append(v)
     return digits
@@ -70,9 +70,9 @@ class BaseCircuit(AbstractCircuit):
     is_mps = False
 
     @staticmethod
-    def all_zero_nodes(n: int, d: int = 2, prefix: str = "qb-") -> List[tn.Node]:
-        prefix = "qd-" if d > 2 else prefix
-        l = [0.0 for _ in range(d)]
+    def all_zero_nodes(n: int, dim: int = 2, prefix: str = "qb-") -> List[tn.Node]:
+        prefix = "qd-" if dim > 2 else prefix
+        l = [0.0 for _ in range(dim)]
         l[0] = 1.0
         nodes = [
             tn.Node(
@@ -541,7 +541,7 @@ def _basis_nod(_k: int) -> Tensor:
         if self.is_dm:
             msconj = []
         if isinstance(l, str):
-            symbols = _decode_basis_label(l, d=self._d, n=self._nqubits)
+            symbols = _decode_basis_label(l, dim=self._d, n=self._nqubits)
             for k in symbols:
                 n = _basis_nod(k)
                 ms.append(tn.Node(n))
@@ -704,7 +704,7 @@ def perfect_sampling(key: Any) -> Any:
                 return r
             r = backend.stack([ri[0] for ri in r])  # type: ignore
             r = backend.cast(r, "int32")
-            ch = sample_bin2int(r, self._nqubits, d=self._d)
+            ch = sample_bin2int(r, self._nqubits, dim=self._d)
         else:  # allow_state
             if batch is None:
                 nbatch = 1
@@ -729,7 +729,7 @@ def perfect_sampling(key: Any) -> Any:
             #     2,
             # )
             if format is None:  # for backward compatibility
-                confg = sample_int2bin(ch, self._nqubits, d=self._d)
+                confg = sample_int2bin(ch, self._nqubits, dim=self._d)
                 prob = backend.gather1d(p, ch)
                 r = list(zip(confg, prob))  # type: ignore
                 if batch is None:
@@ -738,7 +738,7 @@ def perfect_sampling(key: Any) -> Any:
         if self._nqubits > 35:
             jittable = False
         return sample2all(
-            sample=ch, n=self._nqubits, format=format, jittable=jittable, d=self._d
+            sample=ch, n=self._nqubits, format=format, jittable=jittable, dim=self._d
         )
 
     def sample_expectation_ps(
diff --git a/tensorcircuit/circuit.py b/tensorcircuit/circuit.py
index 9220929f..d302d6d9 100644
--- a/tensorcircuit/circuit.py
+++ b/tensorcircuit/circuit.py
@@ -1,6 +1,6 @@
 """
 Quantum circuit: the state simulator.
-Supports qubit (d=2) and qudit (3 <= d <= 36) systems.
+Supports qubit (dim=2) and qudit (3 <= dim <= 36) systems.
  For string-encoded samples/counts, digits use 0–9A–Z where A=10, …, Z=35.
 """
 
@@ -87,7 +87,7 @@ def __init__(
             "split": split,
         }
         if (inputs is None) and (mps_inputs is None):
-            nodes = self.all_zero_nodes(nqubits, d=self._d)
+            nodes = self.all_zero_nodes(nqubits, dim=self._d)
             self._front = [n.get_edge(0) for n in nodes]
         elif inputs is not None:  # provide input function
             inputs = backend.convert_to_tensor(inputs)
@@ -922,7 +922,7 @@ def expectation(
 def expectation(
     *ops: Tuple[tn.Node, List[int]],
     ket: Tensor,
-    d: Optional[int] = None,
+    dim: Optional[int] = None,
     bra: Optional[Tensor] = None,
     conj: bool = True,
     normalization: bool = False,
@@ -974,8 +974,8 @@ def expectation(
     :type ket: Tensor
     :param bra: :math:`bra`, defaults to None, which is the same as ``ket``.
     :type bra: Optional[Tensor], optional
-    :param d: dimension of the circuit (defaults to 2)
-    :type d: int, optional
+    :param dim: dimension of the circuit (defaults to 2)
+    :type dim: int, optional
     :param conj: :math:`bra` changes to the adjoint matrix of :math:`bra`, defaults to True.
     :type conj: bool, optional
     :param normalization: Normalize the :math:`ket` and :math:`bra`, defaults to False.
@@ -984,7 +984,7 @@ def expectation(
     :return: The result of :math:`\\langle bra\\vert ops \\vert ket\\rangle`.
     :rtype: Tensor
     """
-    d = 2 if d is None else d
+    dim = 2 if dim is None else dim
     if bra is None:
         bra = ket
     if isinstance(ket, QuOperator):
@@ -998,7 +998,7 @@ def expectation(
         for op, index in ops:
             if not isinstance(op, tn.Node):
                 # op is only a matrix
-                op = backend.reshaped(op, d)
+                op = backend.reshaped(op, dim)
                 op = gates.Gate(op)
             if isinstance(index, int):
                 index = [index]
@@ -1022,8 +1022,8 @@ def expectation(
         if conj is True:
             bra = backend.conj(bra)
         ket = backend.reshape(ket, [-1])
-        ket = backend.reshaped(ket, d)
-        bra = backend.reshaped(bra, d)
+        ket = backend.reshaped(ket, dim)
+        bra = backend.reshaped(bra, dim)
         n = len(backend.shape_tuple(ket))
         ket = Gate(ket)
         bra = Gate(bra)
@@ -1035,7 +1035,7 @@ def expectation(
         for op, index in ops:
             if not isinstance(op, tn.Node):
                 # op is only a matrix
-                op = backend.reshaped(op, d)
+                op = backend.reshaped(op, dim)
                 op = gates.Gate(op)
             if isinstance(index, int):
                 index = [index]
diff --git a/tensorcircuit/mpscircuit.py b/tensorcircuit/mpscircuit.py
index b8c511ef..d2aad9e5 100644
--- a/tensorcircuit/mpscircuit.py
+++ b/tensorcircuit/mpscircuit.py
@@ -401,8 +401,8 @@ def gate_to_MPO(
 
         nindex = len(index)
         in_dims = tuple(backend.shape_tuple(gate))[:nindex]
-        d = int(in_dims[0])
-        dim_phys_mpo = d * d
+        dim = int(in_dims[0])
+        dim_phys_mpo = dim * dim
 
         order = tuple(np.arange(2 * nindex).reshape(2, nindex).T.flatten().tolist())
         gate = backend.transpose(gate, order)
@@ -429,13 +429,13 @@ def gate_to_MPO(
             if previous_i is not None:
                 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 * d, dtype=backend.dtype(tensors[-1]))
-                    I4 = backend.reshape(eye2d, (bond_dim, d, bond_dim, d))
+                    eye2d = backend.eye(bond_dim * dim, dtype=backend.dtype(tensors[-1]))
+                    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), d, d, int(nright)))
+            tensor = backend.reshape(main_tensor, (int(nleft), dim, dim, int(nright)))
             tensors.append(tensor)
             previous_i = int(i)
 
diff --git a/tensorcircuit/quantum.py b/tensorcircuit/quantum.py
index 9a98f5b7..492854fc 100644
--- a/tensorcircuit/quantum.py
+++ b/tensorcircuit/quantum.py
@@ -57,27 +57,27 @@ def get_all_nodes(edges: Iterable[Edge]) -> List[Node]:
     return nodes
 
 
-def _infer_num_sites(D: int, d: int) -> int:
+def _infer_num_sites(D: int, dim: int) -> int:
     """
     Infer the number of sites (n) from a Hilbert space dimension D
     and local dimension d, assuming D = d**n.
 
     :param D: total Hilbert space dimension (int)
-    :param d: local dimension per site (int)
+    :param dim: local dimension per site (int)
     :return: n such that D == d**n
     :raises ValueError: if D is not an exact power of d
     """
     if not (isinstance(D, int) and D > 0):
         raise ValueError(f"D must be a positive integer, got {D}")
-    if not (isinstance(d, int) and d >= 2):
-        raise ValueError(f"d must be an integer >= 2, got {d}")
+    if not (isinstance(dim, int) and dim >= 2):
+        raise ValueError(f"d must be an integer >= 2, got {dim}")
 
     tmp, n = D, 0
-    while tmp % d == 0 and tmp > 1:
-        tmp //= d
+    while tmp % dim == 0 and tmp > 1:
+        tmp //= dim
         n += 1
     if tmp != 1:
-        raise ValueError(f"Dimension {D} is not a power of local dim {d}")
+        raise ValueError(f"Dimension {D} is not a power of local dim {dim}")
     return n
 
 
@@ -2178,7 +2178,7 @@ def reduced_wavefunction(
     state: Tensor,
     cut: List[int],
     measure: Optional[List[int]] = None,
-    d: Optional[int] = None,
+    dim: Optional[int] = None,
 ) -> Tensor:
     """
     Compute the reduced wavefunction from the quantum state ``state``.
@@ -2193,19 +2193,19 @@ def reduced_wavefunction(
     :type measure: List[int]
     :return: _description_
     :rtype: Tensor
-    :param d: dimension of qudit system
-    :type d: int
+    :param dim: dimension of qudit system
+    :type dim: int
     """
-    d = 2 if d is None else d
+    dim = 2 if dim is None else dim
     if measure is None:
         measure = [0 for _ in cut]
-    s = backend.reshaped(state, d)
+    s = backend.reshaped(state, dim)
     n = len(backend.shape_tuple(s))
     s_node = Gate(s)
     end_nodes = []
     for c, m in zip(cut, measure):
         oh = backend.cast(
-            backend.one_hot(backend.cast(backend.convert_to_tensor(m), "int32"), d),
+            backend.one_hot(backend.cast(backend.convert_to_tensor(m), "int32"), dim),
             dtypestr,
         )
         end_node = Gate(backend.convert_to_tensor(oh))
@@ -2223,7 +2223,7 @@ def reduced_density_matrix(
     cut: Union[int, List[int]],
     p: Optional[Tensor] = None,
     normalize: bool = True,
-    d: Optional[int] = None,
+    dim: Optional[int] = None,
 ) -> Union[Tensor, QuOperator]:
     r"""
     Compute the reduced density matrix from the quantum state ``state``.
@@ -2239,10 +2239,10 @@ def reduced_density_matrix(
     :rtype: Union[Tensor, QuOperator]
     :param normalize: if True, returns a trace 1 density matrix. Otherwise does not normalize.
     :type normalize: bool
-    :param d: dimension of qudit system
-    :type d: int
+    :param dim: dimension of qudit system
+    :type dim: int
     """
-    d = 2 if d is None else d
+    dim = 2 if dim is None else dim
     if isinstance(cut, list) or isinstance(cut, tuple) or isinstance(cut, set):
         traceout = list(cut)
     else:
@@ -2255,19 +2255,19 @@ def reduced_density_matrix(
         return state.partial_trace(traceout)
     if len(state.shape) == 2 and state.shape[0] == state.shape[1]:
         # density operator
-        freedom = _infer_num_sites(state.shape[0], d)
+        freedom = _infer_num_sites(state.shape[0], dim)
         left = traceout + [i for i in range(freedom) if i not in traceout]
         right = [i + freedom for i in left]
 
-        rho = backend.reshape(state, [d] * (2 * freedom))
+        rho = backend.reshape(state, [dim] * (2 * freedom))
         rho = backend.transpose(rho, perm=left + right)
         rho = backend.reshape(
             rho,
             [
-                d ** len(traceout),
-                d ** (freedom - len(traceout)),
-                d ** len(traceout),
-                d ** (freedom - len(traceout)),
+                dim ** len(traceout),
+                dim ** (freedom - len(traceout)),
+                dim ** len(traceout),
+                dim ** (freedom - len(traceout)),
             ],
         )
         if p is None:
@@ -2280,7 +2280,7 @@ def reduced_density_matrix(
             p = backend.reshape(p, [-1])
             rho = backend.einsum("a,aiaj->ij", p, rho)
         rho = backend.reshape(
-            rho, [d ** (freedom - len(traceout)), d ** (freedom - len(traceout))]
+            rho, [dim ** (freedom - len(traceout)), dim ** (freedom - len(traceout))]
         )
         if normalize:
             rho /= backend.trace(rho)
@@ -2288,12 +2288,12 @@ def reduced_density_matrix(
     else:
         w = state / backend.norm(state)
         size = int(backend.sizen(state))
-        freedom = _infer_num_sites(size, d)
+        freedom = _infer_num_sites(size, dim)
         perm = [i for i in range(freedom) if i not in traceout]
         perm = perm + traceout
-        w = backend.reshape(w, [d for _ in range(freedom)])
+        w = backend.reshape(w, [dim for _ in range(freedom)])
         w = backend.transpose(w, perm=perm)
-        w = backend.reshape(w, [-1, d ** len(traceout)])
+        w = backend.reshape(w, [-1, dim ** len(traceout)])
         if p is None:
             rho = w @ backend.adjoint(w)
         else:
@@ -2437,7 +2437,7 @@ def truncated_free_energy(
 
 @op2tensor
 def partial_transpose(
-    rho: Tensor, transposed_sites: List[int], d: Optional[int] = None
+    rho: Tensor, transposed_sites: List[int], dim: Optional[int] = None
 ) -> Tensor:
     """
     _summary_
@@ -2446,13 +2446,13 @@ def partial_transpose(
     :type rho: Tensor
     :param transposed_sites: sites int list to be transposed
     :type transposed_sites: List[int]
-    :param d: dimension of qudit system
-    :type d: int
+    :param dim: dimension of qudit system
+    :type dim: int
     :return: _description_
     :rtype: Tensor
     """
-    d = 2 if d is None else d
-    rho = backend.reshaped(rho, d)
+    dim = 2 if dim is None else dim
+    rho = backend.reshaped(rho, dim)
     rho_node = Gate(rho)
     n = len(rho.shape) // 2
     left_edges = []
@@ -2471,7 +2471,7 @@ def partial_transpose(
 
 @op2tensor
 def entanglement_negativity(
-    rho: Tensor, transposed_sites: List[int], d: Optional[int] = None
+    rho: Tensor, transposed_sites: List[int], dim: Optional[int] = None
 ) -> Tensor:
     """
     _summary_
@@ -2480,12 +2480,12 @@ def entanglement_negativity(
     :type rho: Tensor
     :param transposed_sites: _description_
     :type transposed_sites: List[int]
-    :param d: dimension of qudit system
-    :type d: int
+    :param dim: dimension of qudit system
+    :type dim: int
     :return: _description_
     :rtype: Tensor
     """
-    rhot = partial_transpose(rho, transposed_sites, d=d)
+    rhot = partial_transpose(rho, transposed_sites, dim=dim)
     es = backend.eigvalsh(rhot)
     rhot_m = backend.sum(backend.abs(es))
     return (rhot_m - 1.0) / 2.0
@@ -2493,7 +2493,7 @@ def entanglement_negativity(
 
 @op2tensor
 def log_negativity(
-    rho: Tensor, transposed_sites: List[int], base: str = "e", d: Optional[int] = None
+    rho: Tensor, transposed_sites: List[int], base: str = "e", dim: Optional[int] = None
 ) -> Tensor:
     """
     _summary_
@@ -2504,13 +2504,13 @@ def log_negativity(
     :type transposed_sites: List[int]
     :param base: whether use 2 based log or e based log, defaults to "e"
     :type base: str, optional
-    :param d: dimension of qudit system
-    :type d: int
+    :param dim: dimension of qudit system
+    :type dim: int
     :return: _description_
     :rtype: Tensor
     """
-    d = 2 if d is None else d
-    rhot = partial_transpose(rho, transposed_sites, d)
+    dim = 2 if dim is None else dim
+    rhot = partial_transpose(rho, transposed_sites, dim)
     es = backend.eigvalsh(rhot)
     rhot_m = backend.sum(backend.abs(es))
     een = backend.log(rhot_m)
@@ -2597,7 +2597,7 @@ def double_state(h: Tensor, beta: float = 1) -> Tensor:
 
 @op2tensor
 def mutual_information(
-    s: Tensor, cut: Union[int, List[int]], d: Optional[int] = None
+    s: Tensor, cut: Union[int, List[int]], dim: Optional[int] = None
 ) -> Tensor:
     """
     Mutual information between AB subsystem described by ``cut``.
@@ -2606,12 +2606,12 @@ def mutual_information(
     :type s: Tensor
     :param cut: The AB subsystem.
     :type cut: Union[int, List[int]]
-    :param d: The diagonal matrix in form of Tensor.
-    :type d: Tensor
+    :param dim: The diagonal matrix in form of Tensor.
+    :type dim: Tensor
     :return: The mutual information between AB subsystem described by ``cut``.
     :rtype: Tensor
     """
-    d = 2 if d is None else d
+    dim = 2 if dim is None else dim
     if isinstance(cut, list) or isinstance(cut, tuple) or isinstance(cut, set):
         traceout = list(cut)
     else:
@@ -2619,22 +2619,22 @@ def mutual_information(
 
     if len(s.shape) == 2 and s.shape[0] == s.shape[1]:
         # mixed state
-        n = _infer_num_sites(s.shape[0], d=d)
+        n = _infer_num_sites(s.shape[0], dim=dim)
         hab = entropy(s)
 
         # subsystem a
-        rhoa = reduced_density_matrix(s, traceout, d=d)
+        rhoa = reduced_density_matrix(s, traceout, dim=dim)
         ha = entropy(rhoa)
 
         # need subsystem b as well
         other = tuple(i for i in range(n) if i not in traceout)
-        rhob = reduced_density_matrix(s, other, d=d)  # type: ignore
+        rhob = reduced_density_matrix(s, other, dim=dim)  # type: ignore
         hb = entropy(rhob)
 
     # pure system
     else:
         hab = 0.0
-        rhoa = reduced_density_matrix(s, traceout, d=d)
+        rhoa = reduced_density_matrix(s, traceout, dim=dim)
         ha = hb = entropy(rhoa)
 
     return ha + hb - hab
@@ -2643,7 +2643,7 @@ def mutual_information(
 # measurement results and transformations and correlations below
 
 
-def count_s2d(srepr: Tuple[Tensor, Tensor], n: int, d: Optional[int] = None) -> Tensor:
+def count_s2d(srepr: Tuple[Tensor, Tensor], n: int, dim: Optional[int] = None) -> Tensor:
     """
     measurement shots results, sparse tuple representation to dense representation
     count_vector to count_tuple
@@ -2652,14 +2652,14 @@ def count_s2d(srepr: Tuple[Tensor, Tensor], n: int, d: Optional[int] = None) ->
     :type srepr: Tuple[Tensor, Tensor]
     :param n: number of qubits
     :type n: int
-    :param d: [description], defaults to None
-    :type d: int, optional
+    :param dim: [description], defaults to None
+    :type dim: int, optional
     :return: [description]
     :rtype: Tensor
     """
-    d = 2 if d is None else d
+    dim = 2 if dim is None else dim
     return backend.scatter(
-        backend.cast(backend.zeros([d**n]), srepr[1].dtype),
+        backend.cast(backend.zeros([dim**n]), srepr[1].dtype),
         backend.reshape(srepr[0], [-1, 1]),
         srepr[1],
     )
@@ -2702,7 +2702,7 @@ def count_d2s(drepr: Tensor, eps: float = 1e-7) -> Tuple[Tensor, Tensor]:
 count_t2v = count_d2s
 
 
-def sample_int2bin(sample: Tensor, n: int, d: Optional[int] = None) -> Tensor:
+def sample_int2bin(sample: Tensor, n: int, dim: Optional[int] = None) -> Tensor:
     """
     Convert linear-index samples to per-site digits (base-d).
 
@@ -2710,28 +2710,28 @@ def sample_int2bin(sample: Tensor, n: int, d: Optional[int] = None) -> Tensor:
     :type sample: Tensor
     :param n: number of sites
     :type n: int
-    :param d: local dimension, defaults to 2
-    :type d: int, optional
+    :param dim: local dimension, defaults to 2
+    :type dim: int, optional
     :return: shape [trials, n], entries in [0, d-1]
     :rtype: Tensor
     """
-    d = 2 if d is None else d
-    if d == 2:
+    dim = 2 if dim is None else dim
+    if dim == 2:
         return backend.mod(
             backend.right_shift(sample[..., None], backend.reverse(backend.arange(n))),
             2,
         )
     else:
         pos = backend.reverse(backend.arange(n))
-        base = backend.power(d, pos)
+        base = backend.power(dim, pos)
         digits = backend.mod(
             backend.floor_divide(sample[..., None], base),  # ⌊sample / d**pos⌋
-            d,
+            dim,
         )
         return backend.cast(digits, "int32")
 
 
-def sample_bin2int(sample: Tensor, n: int, d: Optional[int] = None) -> Tensor:
+def sample_bin2int(sample: Tensor, n: int, dim: Optional[int] = None) -> Tensor:
     """
     bin sample to int sample
 
@@ -2739,11 +2739,13 @@ def sample_bin2int(sample: Tensor, n: int, d: Optional[int] = None) -> Tensor:
     :type sample: Tensor
     :param n: number of qubits
     :type n: int
+    :param dim: local dimension, defaults to 2
+    :type dim: int, optional
     :return: in shape [trials]
     :rtype: Tensor
     """
-    d = 2 if d is None else d
-    power = backend.convert_to_tensor([d**j for j in reversed(range(n))])
+    dim = 2 if dim is None else dim
+    power = backend.convert_to_tensor([dim**j for j in reversed(range(n))])
     return backend.sum(sample * power, axis=-1)
 
 
@@ -2751,7 +2753,7 @@ def sample2count(
     sample: Tensor,
     n: int,
     jittable: bool = True,
-    d: Optional[int] = None,
+    dim: Optional[int] = None,
 ) -> Tuple[Tensor, Tensor]:
     """
     sample_int to count_tuple (indices, counts), size = d**n
@@ -2759,11 +2761,11 @@ def sample2count(
     :param sample: linear-index samples, shape [shots]
     :param n: number of sites
     :param jittable: whether to return fixed-size outputs (backend dependent)
-    :param d: local dimension per site, default 2 (qubit)
+    :param dim: local dimension per site, default 2 (qubit)
     :return: (unique_indices, counts)
     """
-    d = 2 if d is None else d
-    size = d**n
+    dim = 2 if dim is None else dim
+    size = dim**n
     if not jittable:
         results = backend.unique_with_counts(sample)  # non-jittable
     else:  # jax specified / fixed-size
@@ -2772,7 +2774,7 @@ def sample2count(
 
 
 def count_vector2dict(
-    count: Tensor, n: int, key: str = "bin", d: Optional[int] = None
+    count: Tensor, n: int, key: str = "bin", dim: Optional[int] = None
 ) -> Dict[Any, int]:
     """
     Convert count_vector to count_dict_bin or count_dict_int.
@@ -2784,28 +2786,28 @@ def count_vector2dict(
     :type n: int
     :param key: can be "int" or "bin", defaults to "bin"
     :type key: str, optional
-    :param d: local dimension (default 2)
-    :type d: int, optional
+    :param dim: local dimension (default 2)
+    :type dim: int, optional
     :return: mapping from configuration to count
     :rtype: Dict[Any, int]
     """
     from .interfaces import which_backend
 
-    d = 2 if d is None else d
+    dim = 2 if dim is None else dim
     b = which_backend(count)
-    out_int = {i: b.numpy(count[i]).item() for i in range(d**n)}
+    out_int = {i: b.numpy(count[i]).item() for i in range(dim**n)}
     if key == "int":
         return out_int
     else:
         out_str = {}
         for k, v in out_int.items():
-            kn = np.base_repr(k, base=d).zfill(n)
+            kn = np.base_repr(k, base=dim).zfill(n)
             out_str[kn] = v
         return out_str
 
 
 def count_tuple2dict(
-    count: Tuple[Tensor, Tensor], n: int, key: str = "bin", d: Optional[int] = None
+    count: Tuple[Tensor, Tensor], n: int, key: str = "bin", dim: Optional[int] = None
 ) -> Dict[Any, int]:
     """
     count_tuple to count_dict_bin or count_dict_int
@@ -2816,12 +2818,12 @@ def count_tuple2dict(
     :type n: int
     :param key: can be "int" or "bin", defaults to "bin"
     :type key: str, optional
-    :param d: local dimension, defaults to 2
-    :type d: int, optional
+    :param dim: local dimension, defaults to 2
+    :type dim: int, optional
     :return: count_dict
     :rtype: Dict[Any, int]
     """
-    d = 2 if d is None else d
+    dim = 2 if dim is None else dim
     out_int = {
         backend.numpy(i).item(): backend.numpy(j).item()
         for i, j in zip(count[0], count[1])
@@ -2832,7 +2834,7 @@ def count_tuple2dict(
     else:
         out_str = {}
         for k, v in out_int.items():
-            kn = np.base_repr(k, base=d).zfill(n)
+            kn = np.base_repr(k, base=dim).zfill(n)
             out_str[kn] = v
         return out_str
 
@@ -2846,7 +2848,7 @@ def measurement_counts(
     random_generator: Optional[Any] = None,
     status: Optional[Tensor] = None,
     jittable: bool = False,
-    d: Optional[int] = None,
+    dim: Optional[int] = None,
 ) -> Any:
     r"""
     Simulate the measuring of each qubit of ``p`` in the computational basis,
@@ -2916,7 +2918,7 @@ def measurement_counts(
             pi = backend.real(backend.conj(state) * state)
         pi = backend.reshape(pi, [-1])
 
-    local_d = 2 if d is None else d
+    local_d = 2 if dim is None else dim
     total_dim = int(backend.shape_tuple(pi)[0])
     n = _infer_num_sites(total_dim, local_d)
 
@@ -2926,9 +2928,9 @@ def measurement_counts(
         elif format == "count_tuple":
             return count_d2s(pi)
         elif format == "count_dict_bin":
-            return count_vector2dict(pi, n, key="bin", d=local_d)
+            return count_vector2dict(pi, n, key="bin", dim=local_d)
         elif format == "count_dict_int":
-            return count_vector2dict(pi, n, key="int", d=local_d)
+            return count_vector2dict(pi, n, key="int", dim=local_d)
         else:
             raise ValueError(f"unsupported format {format} for analytical measurement")
     else:
@@ -2941,7 +2943,7 @@ def measurement_counts(
         # raw_counts = backend.stateful_randc(
         # random_generator, a=drange, shape=counts, p=pi
         # )
-        return sample2all(raw_counts, n, format=format, jittable=jittable, d=local_d)
+        return sample2all(raw_counts, n, format=format, jittable=jittable, dim=local_d)
 
 
 measurement_results = measurement_counts
@@ -2953,7 +2955,7 @@ def sample2all(
     n: int,
     format: str = "count_vector",
     jittable: bool = False,
-    d: Optional[int] = None,
+    dim: Optional[int] = None,
 ) -> Any:
     """
     transform ``sample_int`` or ``sample_bin`` results to other forms specified by ``format``
@@ -2966,12 +2968,12 @@ def sample2all(
     :type format: str, optional
     :param jittable: only applicable to count transformation in jax backend, defaults to False
     :type jittable: bool, optional
-    :param d: local dimension (2 for qubit; >2 for qudit), defaults to 2
-    :type d: Optional[int]
+    :param dim: local dimension (2 for qubit; >2 for qudit), defaults to 2
+    :type dim: Optional[int]
     :return: measurement results specified as ``format``
     :rtype: Any
     """
-    d = 2 if d is None else int(d)
+    dim = 2 if dim is None else int(dim)
     if n > 32:
         assert (
             len(backend.shape_tuple(sample)) == 2
@@ -2985,9 +2987,9 @@ def sample2all(
 
     if len(backend.shape_tuple(sample)) == 1:
         sample_int = sample
-        sample_bin = sample_int2bin(sample, n, d=d)
+        sample_bin = sample_int2bin(sample, n, dim=dim)
     elif len(backend.shape_tuple(sample)) == 2:
-        sample_int = sample_bin2int(sample, n, d=d)
+        sample_int = sample_bin2int(sample, n, dim=dim)
         sample_bin = sample
     else:
         raise ValueError("unrecognized tensor shape for sample")
@@ -2997,15 +2999,15 @@ def sample2all(
     elif format == "sample_bin":
         return sample_bin
     else:
-        count_tuple = sample2count(sample_int, n, jittable=jittable, d=d)
+        count_tuple = sample2count(sample_int, n, jittable=jittable, dim=dim)
         if format == "count_tuple":
             return count_tuple
         elif format == "count_vector":
-            return count_s2d(count_tuple, n, d=d)
+            return count_s2d(count_tuple, n, dim=dim)
         elif format == "count_dict_bin":
-            return count_tuple2dict(count_tuple, n, key="bin", d=d)
+            return count_tuple2dict(count_tuple, n, key="bin", dim=dim)
         elif format == "count_dict_int":
-            return count_tuple2dict(count_tuple, n, key="int", d=d)
+            return count_tuple2dict(count_tuple, n, key="int", dim=dim)
         else:
             raise ValueError(
                 f"unsupported format {format} for finite shots measurement"
@@ -3013,7 +3015,7 @@ def sample2all(
 
 
 def spin_by_basis(
-    n: int, m: int, elements: Tuple[int, int] = (1, -1), d: Optional[int] = None
+    n: int, m: int, elements: Tuple[int, int] = (1, -1), dim: Optional[int] = None
 ) -> Tensor:
     """
     Generate all n-bitstrings as an array, each row is a bitstring basis.
@@ -3034,10 +3036,10 @@ def spin_by_basis(
         all bitstring basis.
     :rtype: Tensor
     """
-    d = len(elements) if d is None else d
+    dim = len(elements) if dim is None else dim
 
     col = backend.convert_to_tensor(np.array(elements, dtype=np.int32).reshape(-1, 1))
-    s = backend.tile(backend.cast(col, "int32"), [d**m, int(d ** (n - m - 1))])
+    s = backend.tile(backend.cast(col, "int32"), [dim**m, int(dim ** (n - m - 1))])
     return backend.reshape(s, [-1])
 
 
@@ -3045,7 +3047,7 @@ def correlation_from_samples(
     index: Sequence[int],
     results: Tensor,
     n: int,
-    d: int = 2,
+    dim: int = 2,
     elements: Optional[Sequence[float]] = None,
 ) -> Tensor:
     r"""
@@ -3057,15 +3059,15 @@ def correlation_from_samples(
     :param index: positions in the basis string
     :param results: samples tensor
     :param n: number of sites
-    :param d: local dimension (default 2)
+    :param dim: local dimension (default 2)
     :param elements: optional mapping of length d from outcome {0..d-1} to values s.
                      If None and d==2, defaults to (1, -1) via the original formula.
     :return: correlation estimate (mean over shots)
     """
     if len(backend.shape_tuple(results)) == 1:
-        results = sample_int2bin(results, n, d=d)
+        results = sample_int2bin(results, n, dim=dim)
 
-    if d == 2 and elements is None:
+    if dim == 2 and elements is None:
         svals = 1 - results * 2  # 0->+1, 1->-1
         r = svals[:, index[0]]
         for i in index[1:]:
@@ -3075,11 +3077,11 @@ def correlation_from_samples(
 
     if elements is None:
         raise ValueError(
-            f"correlation_from_samples requires `elements` mapping for d={d}; "
+            f"correlation_from_samples requires `elements` mapping for d={dim}; "
             f"e.g., for qutrit you might pass elements=(1.0,0.0,-1.0)."
         )
-    if len(elements) != d:
-        raise ValueError(f"`elements` length {len(elements)} != d={d}")
+    if len(elements) != dim:
+        raise ValueError(f"`elements` length {len(elements)} != d={dim}")
 
     evec = backend.cast(backend.convert_to_tensor(np.asarray(elements)), rdtypestr)
 
@@ -3097,7 +3099,7 @@ def correlation_from_samples(
 def correlation_from_counts(
     index: Sequence[int],
     results: Tensor,
-    d: Optional[int] = None,
+    dim: Optional[int] = None,
     elements: Optional[Sequence[float]] = None,
 ) -> Tensor:
     r"""
@@ -3106,34 +3108,34 @@ def correlation_from_counts(
 
     :param index: positions in the basis string
     :param results: probability/count vector of shape d**n (will be normalized)
-    :param d: local dimension (default 2)
-    :param elements: optional mapping of length d from digit {0..d-1} to values s.
+    :param dim: local dimension (default 2)
+    :param dim: optional mapping of length d from digit {0..d-1} to values s.
                      If None and d==2, defaults to (1, -1). For d>2, must be provided.
     :return: correlation expectation from counts
     """
-    d = 2 if d is None else int(d)
-    if d != 2:
-        raise NotImplementedError(f"`d={d}` not implemented.")
+    dim = 2 if dim is None else int(dim)
+    if dim != 2:
+        raise NotImplementedError(f"`d={dim}` not implemented.")
 
     results = backend.reshape(results, [-1])
     results = backend.cast(results, rdtypestr)
     results /= backend.sum(results)
 
-    n = _infer_num_sites(int(results.shape[0]), d=d)
+    n = _infer_num_sites(int(results.shape[0]), dim=dim)
 
-    if d == 2 and elements is None:
+    if dim == 2 and elements is None:
         elems = (1, -1)
     else:
-        if elements is None or len(elements) != d:
+        if elements is None or len(elements) != dim:
             raise ValueError(
-                f"`elements` must be provided with length d={d} for qudit; got {elements}."
+                f"`elements` must be provided with length d={dim} for qudit; got {elements}."
             )
         elems = tuple(elements)  # type: ignore
 
     acc = results
     for i in index:
         acc = acc * backend.cast(
-            spin_by_basis(n, int(i), elements=elems, d=d), acc.dtype
+            spin_by_basis(n, int(i), elements=elems, dim=dim), acc.dtype
         )
 
     return backend.sum(acc)
diff --git a/tensorcircuit/results/counts.py b/tensorcircuit/results/counts.py
index a921f209..53c4e7ef 100644
--- a/tensorcircuit/results/counts.py
+++ b/tensorcircuit/results/counts.py
@@ -151,7 +151,7 @@ def vec2count(vec: Tensor, prune: bool = False, dim: Optional[int] = 2) -> ct:
     if isinstance(vec, list):
         vec = np.array(vec)
     n = int(np.log(int(vec.shape[0])) / np.log(dim) + 1e-9)
-    c: ct = count_vector2dict(vec, n, key="bin", d=dim)  # type: ignore
+    c: ct = count_vector2dict(vec, n, key="bin", dim=dim)  # type: ignore
     if prune:
         c = {k: v for k, v in c.items() if np.abs(v) >= 1e-8}
 

From 5b422825184257606dd4b8c063ea313becbb9e9a Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Mon, 1 Sep 2025 13:26:39 +0800
Subject: [PATCH 30/55] black .

---
 tensorcircuit/mpscircuit.py | 4 +++-
 tensorcircuit/quantum.py    | 4 +++-
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/tensorcircuit/mpscircuit.py b/tensorcircuit/mpscircuit.py
index d2aad9e5..4766faa9 100644
--- a/tensorcircuit/mpscircuit.py
+++ b/tensorcircuit/mpscircuit.py
@@ -429,7 +429,9 @@ def gate_to_MPO(
             if previous_i is not None:
                 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]))
+                    eye2d = backend.eye(
+                        bond_dim * dim, dtype=backend.dtype(tensors[-1])
+                    )
                     I4 = backend.reshape(eye2d, (bond_dim, dim, bond_dim, dim))
                     I4 = backend.transpose(I4, (0, 1, 3, 2))
                     tensors.append(I4)
diff --git a/tensorcircuit/quantum.py b/tensorcircuit/quantum.py
index 52c273b6..4a72e1c2 100644
--- a/tensorcircuit/quantum.py
+++ b/tensorcircuit/quantum.py
@@ -2643,7 +2643,9 @@ def mutual_information(
 # measurement results and transformations and correlations below
 
 
-def count_s2d(srepr: Tuple[Tensor, Tensor], n: int, dim: Optional[int] = None) -> Tensor:
+def count_s2d(
+    srepr: Tuple[Tensor, Tensor], n: int, dim: Optional[int] = None
+) -> Tensor:
     """
     measurement shots results, sparse tuple representation to dense representation
     count_vector to count_tuple

From 28af7a28aaa1a5f82357e2a35caa627b261a6475 Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Mon, 1 Sep 2025 13:37:59 +0800
Subject: [PATCH 31/55] a bug fixed

---
 tensorcircuit/results/counts.py | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/tensorcircuit/results/counts.py b/tensorcircuit/results/counts.py
index 53c4e7ef..772d2cee 100644
--- a/tensorcircuit/results/counts.py
+++ b/tensorcircuit/results/counts.py
@@ -90,7 +90,9 @@ def marginal_count(count: ct, keep_list: Sequence[int]) -> ct:
     return reverse_count(ncount)
 
 
-def count2vec(count: ct, normalization: bool = True, dim: Optional[int] = 2) -> Tensor:
+def count2vec(
+    count: ct, normalization: bool = True, dim: Optional[int] = None
+) -> Tensor:
     """
     Convert a dictionary of counts (with string keys) to a probability/count vector.
 
@@ -121,6 +123,8 @@ def parse_key(_k: str) -> List[int]:
     if not count:
         return np.array([], dtype=float)
 
+    dim = 2 if dim is None else dim
+
     n = len(next(iter(count)).upper())
     prob = np.zeros(dim**n, dtype=float)
     shots = float(sum(count.values())) if normalization else 1.0
@@ -136,7 +140,7 @@ def parse_key(_k: str) -> List[int]:
     return prob
 
 
-def vec2count(vec: Tensor, prune: bool = False, dim: Optional[int] = 2) -> ct:
+def vec2count(vec: Tensor, prune: bool = False, dim: Optional[int] = None) -> ct:
     """
     Map a count/probability vector of length D to a dictionary with base-d string keys (0-9A-Z).
     Only generate string keys when d ≤ 36; if d is inferred to be > 36, raise a NotImplementedError.
@@ -146,8 +150,9 @@ def vec2count(vec: Tensor, prune: bool = False, dim: Optional[int] = 2) -> ct:
     :param dim: Dimensionality of the vector, defaults to 2
     :return: {base-d string key: value}, key length n
     """
-    from ..quantum import count_vector2dict, _infer_num_sites
+    from ..quantum import count_vector2dict
 
+    dim = 2 if dim is None else dim
     if isinstance(vec, list):
         vec = np.array(vec)
     n = int(np.log(int(vec.shape[0])) / np.log(dim) + 1e-9)

From 75b6c1d1fcd43d503e9b08752dffbdc1284f29a7 Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Mon, 1 Sep 2025 17:49:42 +0800
Subject: [PATCH 32/55] I found that the correlation in qudit systems is
 ill-defined, so does spin_by_basis, so I cancel the changes in these three
 functions.

---
 tensorcircuit/quantum.py | 130 +++++++++++++--------------------------
 1 file changed, 43 insertions(+), 87 deletions(-)

diff --git a/tensorcircuit/quantum.py b/tensorcircuit/quantum.py
index 4a72e1c2..9e324bfc 100644
--- a/tensorcircuit/quantum.py
+++ b/tensorcircuit/quantum.py
@@ -3016,9 +3016,7 @@ def sample2all(
             )
 
 
-def spin_by_basis(
-    n: int, m: int, elements: Tuple[int, int] = (1, -1), dim: Optional[int] = None
-) -> Tensor:
+def spin_by_basis(n: int, m: int, elements: Tuple[int, int] = (1, -1)) -> Tensor:
     """
     Generate all n-bitstrings as an array, each row is a bitstring basis.
     Return m-th col.
@@ -3038,109 +3036,67 @@ def spin_by_basis(
         all bitstring basis.
     :rtype: Tensor
     """
-    dim = len(elements) if dim is None else dim
-
-    col = backend.convert_to_tensor(np.array(elements, dtype=np.int32).reshape(-1, 1))
-    s = backend.tile(backend.cast(col, "int32"), [dim**m, int(dim ** (n - m - 1))])
+    s = backend.tile(
+        backend.cast(
+            backend.convert_to_tensor(np.array([[elements[0]], [elements[1]]])), "int32"
+        ),
+        [2**m, int(2 ** (n - m - 1))],
+    )
     return backend.reshape(s, [-1])
 
 
-def correlation_from_samples(
-    index: Sequence[int],
-    results: Tensor,
-    n: int,
-    dim: int = 2,
-    elements: Optional[Sequence[float]] = None,
-) -> Tensor:
+def correlation_from_samples(index: Sequence[int], results: Tensor, n: int) -> Tensor:
     r"""
-    Compute :math:`\prod_{i\in \text{index}} s_i` from measurement shots,
-    where each site value :math:`s_i` is mapped from the digit outcome.
+    Compute :math:`\prod_{i\in \\text{index}} s_i (s=\pm 1)`,
+    Results is in the format of "sample_int" or "sample_bin"
 
-    Results can be "sample_int" ([shots]) or "sample_bin" ([shots, n]).
-
-    :param index: positions in the basis string
-    :param results: samples tensor
-    :param n: number of sites
-    :param dim: local dimension (default 2)
-    :param elements: optional mapping of length d from outcome {0..d-1} to values s.
-                     If None and d==2, defaults to (1, -1) via the original formula.
-    :return: correlation estimate (mean over shots)
+    :param index: list of int, indicating the position in the bitstring
+    :type index: Sequence[int]
+    :param results: sample tensor
+    :type results: Tensor
+    :param n: number of qubits
+    :type n: int
+    :return: Correlation expectation from measurement shots
+    :rtype: Tensor
     """
     if len(backend.shape_tuple(results)) == 1:
-        results = sample_int2bin(results, n, dim=dim)
-
-    if dim == 2 and elements is None:
-        svals = 1 - results * 2  # 0->+1, 1->-1
-        r = svals[:, index[0]]
-        for i in index[1:]:
-            r *= svals[:, i]
-        r = backend.cast(r, rdtypestr)
-        return backend.mean(r)
-
-    if elements is None:
-        raise ValueError(
-            f"correlation_from_samples requires `elements` mapping for d={dim}; "
-            f"e.g., for qutrit you might pass elements=(1.0,0.0,-1.0)."
-        )
-    if len(elements) != dim:
-        raise ValueError(f"`elements` length {len(elements)} != d={dim}")
-
-    evec = backend.cast(backend.convert_to_tensor(np.asarray(elements)), rdtypestr)
-
-    col = backend.cast(results[:, index[0]], "int32")
-    r = backend.gather1d(evec, col)  # shape [shots]
-
+        results = sample_int2bin(results, n)
+    results = 1 - results * 2
+    r = results[:, index[0]]
     for i in index[1:]:
-        col = backend.cast(results[:, i], "int32")
-        r *= backend.gather1d(evec, col)
-
+        r *= results[:, i]
     r = backend.cast(r, rdtypestr)
     return backend.mean(r)
 
 
-def correlation_from_counts(
-    index: Sequence[int],
-    results: Tensor,
-    dim: Optional[int] = None,
-    elements: Optional[Sequence[float]] = None,
-) -> Tensor:
+def correlation_from_counts(index: Sequence[int], results: Tensor) -> Tensor:
     r"""
-    Compute :math:`\prod_{i\in \text{index}} s_i` where the probability for each
-    basis label is given by a count/probability vector ``results`` ("count_vector").
+    Compute :math:`\prod_{i\in \\text{index}} s_i`,
+    where the probability for each bitstring is given as a vector ``results``.
+    Results is in the format of "count_vector"
 
-    :param index: positions in the basis string
-    :param results: probability/count vector of shape d**n (will be normalized)
-    :param dim: local dimension (default 2)
-    :param dim: optional mapping of length d from digit {0..d-1} to values s.
-                     If None and d==2, defaults to (1, -1). For d>2, must be provided.
-    :return: correlation expectation from counts
-    """
-    dim = 2 if dim is None else int(dim)
-    if dim != 2:
-        raise NotImplementedError(f"`d={dim}` not implemented.")
+    :Example:
 
+    >>> prob = tc.array_to_tensor(np.array([0.6, 0.4, 0, 0]))
+    >>> qu.correlation_from_counts([0, 1], prob)
+    (0.20000002+0j)
+    >>> qu.correlation_from_counts([1], prob)
+    (0.20000002+0j)
+
+    :param index: list of int, indicating the position in the bitstring
+    :type index: Sequence[int]
+    :param results: probability vector of shape 2^n
+    :type results: Tensor
+    :return: Correlation expectation from measurement shots.
+    :rtype: Tensor
+    """
     results = backend.reshape(results, [-1])
     results = backend.cast(results, rdtypestr)
     results /= backend.sum(results)
-
-    n = _infer_num_sites(int(results.shape[0]), dim=dim)
-
-    if dim == 2 and elements is None:
-        elems = (1, -1)
-    else:
-        if elements is None or len(elements) != dim:
-            raise ValueError(
-                f"`elements` must be provided with length d={dim} for qudit; got {elements}."
-            )
-        elems = tuple(elements)  # type: ignore
-
-    acc = results
+    n = int(np.log(results.shape[0]) / np.log(2))
     for i in index:
-        acc = acc * backend.cast(
-            spin_by_basis(n, int(i), elements=elems, dim=dim), acc.dtype
-        )
-
-    return backend.sum(acc)
+        results = results * backend.cast(spin_by_basis(n, i), results.dtype)
+    return backend.sum(results)
 
 
 # @op2tensor

From 3832c028473c9e6abded69af29b1d0551b4c03ef Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Mon, 1 Sep 2025 18:09:09 +0800
Subject: [PATCH 33/55] I changed this into a more random method.

---
 tensorcircuit/basecircuit.py | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/tensorcircuit/basecircuit.py b/tensorcircuit/basecircuit.py
index 615be25c..faf2ad0e 100644
--- a/tensorcircuit/basecircuit.py
+++ b/tensorcircuit/basecircuit.py
@@ -472,12 +472,10 @@ def measure_jit(
             else:
                 zero_r = backend.cast(backend.convert_to_tensor(0.0), rdtypestr)
                 tiny_r = backend.cast(backend.convert_to_tensor(1e-12), rdtypestr)
-                pu = backend.real(backend.diagonal(rho))
-                pu = backend.clip(pu, zero_r, one_r)
-                pu = pu + tiny_r * (
-                    backend.ones((self._d,), dtype=rdtypestr)
-                    / backend.cast(backend.convert_to_tensor(float(self._d)), rdtypestr)
-                )
+                pu = backend.clip(backend.real(backend.diagonal(rho)), zero_r, one_r)
+                phi = backend.cast(backend.convert_to_tensor((np.sqrt(5.0) - 1.0) / 2.0), rdtypestr)
+                frac = backend.mod(backend.cast(backend.arange(self._d), rdtypestr) * phi, one_r)
+                pu = pu + tiny_r * (frac + tiny_r)
                 pu = pu / backend.sum(pu)
                 cdf = backend.cumsum(pu)
                 if status is None:

From 32893c3a87435a2b07b257cbb51e5395ee63109d Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Mon, 1 Sep 2025 18:09:33 +0800
Subject: [PATCH 34/55] black .

---
 tensorcircuit/basecircuit.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/tensorcircuit/basecircuit.py b/tensorcircuit/basecircuit.py
index faf2ad0e..e8297209 100644
--- a/tensorcircuit/basecircuit.py
+++ b/tensorcircuit/basecircuit.py
@@ -473,8 +473,12 @@ def measure_jit(
                 zero_r = backend.cast(backend.convert_to_tensor(0.0), rdtypestr)
                 tiny_r = backend.cast(backend.convert_to_tensor(1e-12), rdtypestr)
                 pu = backend.clip(backend.real(backend.diagonal(rho)), zero_r, one_r)
-                phi = backend.cast(backend.convert_to_tensor((np.sqrt(5.0) - 1.0) / 2.0), rdtypestr)
-                frac = backend.mod(backend.cast(backend.arange(self._d), rdtypestr) * phi, one_r)
+                phi = backend.cast(
+                    backend.convert_to_tensor((np.sqrt(5.0) - 1.0) / 2.0), rdtypestr
+                )
+                frac = backend.mod(
+                    backend.cast(backend.arange(self._d), rdtypestr) * phi, one_r
+                )
                 pu = pu + tiny_r * (frac + tiny_r)
                 pu = pu / backend.sum(pu)
                 cdf = backend.cumsum(pu)

From 3802c7dcb3fdcee68134a62987f0ef99f0e84d17 Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Mon, 1 Sep 2025 18:10:15 +0800
Subject: [PATCH 35/55] add test for conversions.

---
 tests/test_quantum.py | 119 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 119 insertions(+)

diff --git a/tests/test_quantum.py b/tests/test_quantum.py
index 816f6dba..8cbf13a4 100644
--- a/tests/test_quantum.py
+++ b/tests/test_quantum.py
@@ -1034,3 +1034,122 @@ def test_u1_project(backend):
     s1 = tc.quantum.u1_project(s, 8, 3)
     assert s1.shape[-1] == 56
     np.testing.assert_allclose(tc.quantum.u1_enlarge(s1, 8, 3), s)
+
+
+@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
+def test_sample_int2bin(backend):
+    for n, dim in [(3, 2), (3, 3), (4, 5), (4, 10), (5, 36)]:
+        trials = 10
+        max_val = dim**n
+        samples = np.random.randint(0, max_val, size=trials)
+
+        digits = np.asarray(tc.quantum.sample_int2bin(samples, n=n, dim=dim))
+        print(digits)
+
+        assert digits.shape == (trials, n)
+
+        assert digits.min() >= 0 and digits.max() < dim
+
+        weights = dim ** np.arange(n - 1, -1, -1)
+        recon = (digits * weights).sum(axis=1)
+        assert np.array_equal(recon, samples)
+
+
+@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
+def test_count_s2d(backend):
+    cases = [
+        (3, None, [0, 3, 7], [5, 2, 1]),
+        (2, 3, [0, 1, 8], [1, 2, 3]),
+        (3, 5, [0, 57, 124], [4, 1, 6]),
+        (2, 10, [7, 42, 99], [2, 5, 1]),
+        (2, 36, [0, 35, 1295], [9, 8, 7]),
+    ]
+
+    for n, dim, idx, cnt in cases:
+        D = 2 if dim is None else dim
+        size = D**n
+
+        s_indices = np.asarray(idx, dtype=np.int64)
+        s_counts_i = np.asarray(cnt, dtype=np.int64)
+
+        dense_i = np.asarray(
+            tc.quantum.count_s2d((s_indices, s_counts_i), n=n, dim=dim)
+        )
+        print("int case:", n, dim, dense_i)
+
+        assert dense_i.shape == (size,)
+        expected_i = np.zeros(size, dtype=np.int64)
+        expected_i[s_indices] = s_counts_i
+        np.testing.assert_array_equal(dense_i, expected_i)
+
+        s_counts_f = np.asarray(cnt, dtype=np.float32)
+        dense_f = np.asarray(
+            tc.quantum.count_s2d((s_indices, s_counts_f), n=n, dim=dim)
+        )
+        print("float case:", n, dim, dense_f)
+
+        assert dense_f.shape == (size,)
+        assert dense_f.dtype == np.float32
+        expected_f = np.zeros(size, dtype=np.float32)
+        expected_f[s_indices] = s_counts_f
+        np.testing.assert_array_equal(dense_f, expected_f)
+
+
+@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
+def test_sample_bin2int(backend):
+    for n, dim in [(3, 2), (3, 3), (4, 5), (4, 10), (5, 36)]:
+        trials = 10
+        digits = np.random.randint(0, dim, size=(trials, n))
+        ints = np.asarray(tc.quantum.sample_bin2int(digits, n=n, dim=dim))
+
+        assert ints.shape == (trials,)
+        weights = dim ** np.arange(n - 1, -1, -1)
+        expected = (digits * weights).sum(axis=1)
+        np.testing.assert_array_equal(ints, expected)
+
+
+@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
+def test_count_vector2dict(backend):
+    for n, dim in [(3, 2), (3, 3), (4, 5), (3, 10), (3, 16), (2, 36)]:
+        size = dim**n
+        counts_np = np.random.randint(0, 7, size=size, dtype=np.int64)
+        counts = tc.backend.convert_to_tensor(counts_np)
+
+        out_int = tc.quantum.count_vector2dict(counts, n=n, key="int", dim=dim)
+        assert isinstance(out_int, dict)
+        assert len(out_int) == size
+        for i in range(size):
+            assert out_int[i] == int(counts_np[i])
+
+        out_bin = tc.quantum.count_vector2dict(counts, n=n, key="bin", dim=dim)
+        assert isinstance(out_bin, dict)
+        assert len(out_bin) == size
+        expected_keys = [np.base_repr(i, base=dim).zfill(n) for i in range(size)]
+        assert set(out_bin.keys()) == set(expected_keys)
+        for i, k in enumerate(expected_keys):
+            assert out_bin[k] == int(counts_np[i])
+
+
+@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
+def test_count_tuple2dict(backend):
+    for n, dim in [(3, 2), (3, 3), (2, 5), (3, 10), (2, 16), (2, 36)]:
+        size = dim**n
+        indices = np.random.choice(size, size=min(5, size), replace=False)
+        counts = np.random.randint(1, 10, size=len(indices))
+
+        idx_tensor = tc.backend.cast(tc.backend.convert_to_tensor(indices), "int64")
+        cnt_tensor = tc.backend.cast(tc.backend.convert_to_tensor(counts), "int64")
+        count_tuple = (idx_tensor, cnt_tensor)
+
+        out_int = tc.quantum.count_tuple2dict(count_tuple, n=n, key="int", dim=dim)
+        assert isinstance(out_int, dict)
+        assert set(out_int.keys()) == set(indices)
+        for i, c in zip(indices, counts):
+            assert out_int[int(i)] == int(c)
+
+        out_bin = tc.quantum.count_tuple2dict(count_tuple, n=n, key="bin", dim=dim)
+        assert isinstance(out_bin, dict)
+        expected_keys = [np.base_repr(int(i), base=dim).zfill(n) for i in indices]
+        assert set(out_bin.keys()) == set(expected_keys)
+        for k, c in zip(expected_keys, counts):
+            assert out_bin[k] == int(c)

From 8b41e593a2277d025d9595be86d0b09eec32e2f2 Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Mon, 1 Sep 2025 19:06:26 +0800
Subject: [PATCH 36/55] formatted code

---
 tensorcircuit/circuit.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorcircuit/circuit.py b/tensorcircuit/circuit.py
index d302d6d9..5e63f713 100644
--- a/tensorcircuit/circuit.py
+++ b/tensorcircuit/circuit.py
@@ -1,7 +1,7 @@
 """
 Quantum circuit: the state simulator.
 Supports qubit (dim=2) and qudit (3 <= dim <= 36) systems.
- For string-encoded samples/counts, digits use 0–9A–Z where A=10, …, Z=35.
+For string-encoded samples/counts, digits use 0–9A–Z where A=10, …, Z=35.
 """
 
 # pylint: disable=invalid-name

From 762459654756d9ff7f2ed3ff405f89f3b27fa310 Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Mon, 1 Sep 2025 20:12:52 +0800
Subject: [PATCH 37/55] remove redundant code.

---
 tensorcircuit/abstractcircuit.py |  1 -
 tensorcircuit/circuit.py         | 13 +++----------
 tensorcircuit/densitymatrix.py   |  9 +--------
 tensorcircuit/mpscircuit.py      | 10 ++--------
 4 files changed, 6 insertions(+), 27 deletions(-)

diff --git a/tensorcircuit/abstractcircuit.py b/tensorcircuit/abstractcircuit.py
index 1e8594c6..6dd6d07d 100644
--- a/tensorcircuit/abstractcircuit.py
+++ b/tensorcircuit/abstractcircuit.py
@@ -69,7 +69,6 @@
 class AbstractCircuit:
     _nqubits: int
     _d: int = 2
-    _qudit: bool = False
     _qir: List[Dict[str, Any]]
     _extra_qir: List[Dict[str, Any]]
     inputs: Tensor
diff --git a/tensorcircuit/circuit.py b/tensorcircuit/circuit.py
index 5e63f713..be5d8da3 100644
--- a/tensorcircuit/circuit.py
+++ b/tensorcircuit/circuit.py
@@ -48,10 +48,10 @@ def __init__(
         inputs: Optional[Tensor] = None,
         mps_inputs: Optional[QuOperator] = None,
         split: Optional[Dict[str, Any]] = None,
-        **kwargs: Any,
     ) -> None:
         r"""
         Circuit object based on state simulator.
+        Do not use this class with d!=2 directly, use tc.QuditCircuit instead for qudit systems.
 
         :param nqubits: The number of qubits in the circuit.
         :type nqubits: int
@@ -67,13 +67,6 @@ def __init__(
         :type split: Optional[Dict[str, Any]]
         """
         self._d = 2 if dim is None else dim
-        self._qudit: bool = kwargs.get("qudit", False)
-        if not self._qudit and self._d != 2:
-            raise ValueError(
-                f"Circuit only supports qubits (dim=2). "
-                f"You passed dim={self._d}. Please use `QuditCircuit` instead."
-            )
-
         self.inputs = inputs
         self.mps_inputs = mps_inputs
         self.split = split
@@ -740,7 +733,7 @@ def get_quoperator(self) -> QuOperator:
         :rtype: QuOperator
         """
         mps = identity([self._d for _ in range(self._nqubits)])
-        c = Circuit(self._nqubits, self._d, qudit=self._qudit)
+        c = Circuit(self._nqubits, self._d)
         ns, es = self._copy()
         c._nodes = ns
         c._front = es
@@ -761,7 +754,7 @@ def matrix(self) -> Tensor:
         :rtype: Tensor
         """
         mps = identity([self._d for _ in range(self._nqubits)])
-        c = Circuit(self._nqubits, self._d, qudit=self._qudit)
+        c = Circuit(self._nqubits, self._d)
         ns, es = self._copy()
         c._nodes = ns
         c._front = es
diff --git a/tensorcircuit/densitymatrix.py b/tensorcircuit/densitymatrix.py
index 514b5264..06421daf 100644
--- a/tensorcircuit/densitymatrix.py
+++ b/tensorcircuit/densitymatrix.py
@@ -36,10 +36,10 @@ def __init__(
         dminputs: Optional[Tensor] = None,
         mpo_dminputs: Optional[QuOperator] = None,
         split: Optional[Dict[str, Any]] = None,
-        **kwargs: Any,
     ) -> None:
         """
         The density matrix simulator based on tensornetwork engine.
+        Do not use this class with d!=2 directly
 
         :param nqubits: Number of qubits
         :type nqubits: int
@@ -58,13 +58,6 @@ def __init__(
         :type split: Optional[Dict[str, Any]]
         """
         self._d = 2 if dim is None else dim
-        self._qudit: bool = kwargs.get("qudit", False)
-        if not self._qudit and self._d != 2:
-            raise ValueError(
-                f"Circuit only supports qubits (dim=2). "
-                f"You passed dim={self._d}. Please use `QuditCircuit` instead."
-            )
-
         if not empty:
             if (
                 (inputs is None)
diff --git a/tensorcircuit/mpscircuit.py b/tensorcircuit/mpscircuit.py
index 4766faa9..977ad37a 100644
--- a/tensorcircuit/mpscircuit.py
+++ b/tensorcircuit/mpscircuit.py
@@ -94,10 +94,10 @@ def __init__(
         tensors: Optional[Sequence[Tensor]] = None,
         wavefunction: Optional[Union[QuVector, Tensor]] = None,
         split: Optional[Dict[str, Any]] = None,
-        **kwargs: Any,
     ) -> 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
@@ -115,13 +115,7 @@ def __init__(
         :param split: Split rules
         :type split: Any
         """
-        self._qudit: bool = kwargs.get("qudit", False)
-        if not self._qudit and self._d != 2:
-            raise ValueError(
-                f"Circuit only supports qubits (dim=2). "
-                f"You passed dim={self._d}. Please use `QuditCircuit` instead."
-            )
-
+        self._d = 2 if dim is None else dim
         self.circuit_param = {
             "nqubits": nqubits,
             "center_position": center_position,

From 5c24c4e7bcc0f6fef5bd0139bec83deb34c2a04d Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Tue, 2 Sep 2025 09:08:57 +0800
Subject: [PATCH 38/55] set parameter `dim` be the last args

---
 tensorcircuit/basecircuit.py   | 6 +++---
 tensorcircuit/circuit.py       | 6 +++---
 tensorcircuit/densitymatrix.py | 2 +-
 tensorcircuit/mpscircuit.py    | 4 ++--
 4 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/tensorcircuit/basecircuit.py b/tensorcircuit/basecircuit.py
index e8297209..90a3917e 100644
--- a/tensorcircuit/basecircuit.py
+++ b/tensorcircuit/basecircuit.py
@@ -36,7 +36,7 @@
 Tensor = Any
 
 
-def _decode_basis_label(label: str, dim: int, n: int) -> List[int]:
+def _decode_basis_label(label: str, n: int, dim: int) -> List[int]:
     if dim > 36:
         raise NotImplementedError(
             f"String basis label supports d<=36 (0–9A–Z). Got dim={dim}. "
@@ -69,7 +69,7 @@ class BaseCircuit(AbstractCircuit):
     is_mps = False
 
     @staticmethod
-    def all_zero_nodes(n: int, dim: int = 2, prefix: str = "qb-") -> List[tn.Node]:
+    def all_zero_nodes(n: int, prefix: str = "qb-", dim: int = 2) -> List[tn.Node]:
         prefix = "qd-" if dim > 2 else prefix
         l = [0.0 for _ in range(dim)]
         l[0] = 1.0
@@ -542,7 +542,7 @@ def _basis_nod(_k: int) -> Tensor:
         if self.is_dm:
             msconj = []
         if isinstance(l, str):
-            symbols = _decode_basis_label(l, dim=self._d, n=self._nqubits)
+            symbols = _decode_basis_label(l, n=self._nqubits, dim=self._d)
             for k in symbols:
                 n = _basis_nod(k)
                 ms.append(tn.Node(n))
diff --git a/tensorcircuit/circuit.py b/tensorcircuit/circuit.py
index be5d8da3..11fd29ed 100644
--- a/tensorcircuit/circuit.py
+++ b/tensorcircuit/circuit.py
@@ -44,10 +44,10 @@ class Circuit(BaseCircuit):
     def __init__(
         self,
         nqubits: int,
-        dim: Optional[int] = None,
         inputs: Optional[Tensor] = None,
         mps_inputs: Optional[QuOperator] = None,
         split: Optional[Dict[str, Any]] = None,
+        dim: Optional[int] = None,
     ) -> None:
         r"""
         Circuit object based on state simulator.
@@ -74,10 +74,10 @@ def __init__(
 
         self.circuit_param = {
             "nqubits": nqubits,
-            "dim": dim,
             "inputs": inputs,
             "mps_inputs": mps_inputs,
             "split": split,
+            "dim": dim,
         }
         if (inputs is None) and (mps_inputs is None):
             nodes = self.all_zero_nodes(nqubits, dim=self._d)
@@ -915,10 +915,10 @@ def expectation(
 def expectation(
     *ops: Tuple[tn.Node, List[int]],
     ket: Tensor,
-    dim: Optional[int] = None,
     bra: Optional[Tensor] = None,
     conj: bool = True,
     normalization: bool = False,
+    dim: Optional[int] = None,
 ) -> Tensor:
     """
     Compute :math:`\\langle bra\\vert ops \\vert ket\\rangle`.
diff --git a/tensorcircuit/densitymatrix.py b/tensorcircuit/densitymatrix.py
index 06421daf..690f4e10 100644
--- a/tensorcircuit/densitymatrix.py
+++ b/tensorcircuit/densitymatrix.py
@@ -29,13 +29,13 @@ class DMCircuit(BaseCircuit):
     def __init__(
         self,
         nqubits: int,
-        dim: Optional[int] = None,
         empty: bool = False,
         inputs: Optional[Tensor] = None,
         mps_inputs: Optional[QuOperator] = None,
         dminputs: Optional[Tensor] = None,
         mpo_dminputs: Optional[QuOperator] = None,
         split: Optional[Dict[str, Any]] = None,
+        dim: Optional[int] = None,
     ) -> None:
         """
         The density matrix simulator based on tensornetwork engine.
diff --git a/tensorcircuit/mpscircuit.py b/tensorcircuit/mpscircuit.py
index 977ad37a..c9ec6a9d 100644
--- a/tensorcircuit/mpscircuit.py
+++ b/tensorcircuit/mpscircuit.py
@@ -89,11 +89,11 @@ class MPSCircuit(AbstractCircuit):
     def __init__(
         self,
         nqubits: int,
-        dim: Optional[int] = None,
         center_position: Optional[int] = None,
         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.
@@ -851,7 +851,7 @@ def normalize(self) -> None:
 
     def amplitude(self, l: str) -> Tensor:
         assert len(l) == self._nqubits
-        idx_list = _decode_basis_label(l, self._d, self._nqubits)
+        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]
 

From 53ae6a3a6a638a03bc405ebac703ba6ce8ae41f8 Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Tue, 2 Sep 2025 10:29:43 +0800
Subject: [PATCH 39/55] bug fixed

---
 tensorcircuit/circuit.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorcircuit/circuit.py b/tensorcircuit/circuit.py
index 11fd29ed..0bca86ec 100644
--- a/tensorcircuit/circuit.py
+++ b/tensorcircuit/circuit.py
@@ -733,7 +733,7 @@ def get_quoperator(self) -> QuOperator:
         :rtype: QuOperator
         """
         mps = identity([self._d for _ in range(self._nqubits)])
-        c = Circuit(self._nqubits, self._d)
+        c = Circuit(self._nqubits, dim=self._d)
         ns, es = self._copy()
         c._nodes = ns
         c._front = es
@@ -754,7 +754,7 @@ def matrix(self) -> Tensor:
         :rtype: Tensor
         """
         mps = identity([self._d for _ in range(self._nqubits)])
-        c = Circuit(self._nqubits, self._d)
+        c = Circuit(self._nqubits, dim=self._d)
         ns, es = self._copy()
         c._nodes = ns
         c._front = es

From ffb200637a2dfd6348ab3ea977241dafa4b07744 Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Tue, 2 Sep 2025 10:30:47 +0800
Subject: [PATCH 40/55] change strategy

---
 tensorcircuit/basecircuit.py | 27 ++++++++++-----------------
 1 file changed, 10 insertions(+), 17 deletions(-)

diff --git a/tensorcircuit/basecircuit.py b/tensorcircuit/basecircuit.py
index 90a3917e..7ab5cb7b 100644
--- a/tensorcircuit/basecircuit.py
+++ b/tensorcircuit/basecircuit.py
@@ -471,28 +471,21 @@ def measure_jit(
                 p = p * (pu * (-1) ** sign + sign)
             else:
                 zero_r = backend.cast(backend.convert_to_tensor(0.0), rdtypestr)
-                tiny_r = backend.cast(backend.convert_to_tensor(1e-12), rdtypestr)
                 pu = backend.clip(backend.real(backend.diagonal(rho)), zero_r, one_r)
-                phi = backend.cast(
-                    backend.convert_to_tensor((np.sqrt(5.0) - 1.0) / 2.0), rdtypestr
-                )
-                frac = backend.mod(
-                    backend.cast(backend.arange(self._d), rdtypestr) * phi, one_r
-                )
-                pu = pu + tiny_r * (frac + tiny_r)
                 pu = pu / backend.sum(pu)
-                cdf = backend.cumsum(pu)
                 if status is None:
-                    r = backend.implicit_randu()[0]
-                    r = backend.real(backend.cast(r, rdtypestr))
+                    a = backend.arange(self._d)
+                    k_out = backend.implicit_randc(a=a, shape=())[0]
+                    k_out = backend.cast(k_out, "int32")
                 else:
                     r = backend.real(backend.cast(status[k], rdtypestr))
-                k_out = backend.sum(backend.cast(cdf <= r, "int32"))
-                k_out = backend.clip(
-                    k_out,
-                    backend.cast(backend.convert_to_tensor(0), "int32"),
-                    backend.cast(backend.convert_to_tensor(self._d - 1), "int32"),
-                )
+                    cdf = backend.cumsum(pu)
+                    k_out = backend.sum(backend.cast(cdf < r, "int32"))
+                    k_out = backend.clip(
+                        k_out,
+                        backend.cast(backend.convert_to_tensor(0), "int32"),
+                        backend.cast(backend.convert_to_tensor(self._d - 1), "int32"),
+                    )
                 sample.append(backend.cast(k_out, rdtypestr))
                 p = p * backend.cast(pu[k_out], rdtypestr)
         sample = backend.real(backend.stack(sample))

From f079f24bfc405e0146d102bfc00794dd575c74d2 Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Tue, 2 Sep 2025 13:34:13 +0800
Subject: [PATCH 41/55] optimized

---
 tensorcircuit/basecircuit.py |  7 +++---
 tensorcircuit/mpscircuit.py  | 41 +++++++++++++++++-------------------
 2 files changed, 23 insertions(+), 25 deletions(-)

diff --git a/tensorcircuit/basecircuit.py b/tensorcircuit/basecircuit.py
index 7ab5cb7b..cda2fdb7 100644
--- a/tensorcircuit/basecircuit.py
+++ b/tensorcircuit/basecircuit.py
@@ -474,13 +474,14 @@ def measure_jit(
                 pu = backend.clip(backend.real(backend.diagonal(rho)), zero_r, one_r)
                 pu = pu / backend.sum(pu)
                 if status is None:
-                    a = backend.arange(self._d)
-                    k_out = backend.implicit_randc(a=a, shape=())[0]
+                    k_out = backend.implicit_randc(
+                        self._d, shape=1, p=backend.cast(pu, rdtypestr)
+                    )[0]
                     k_out = backend.cast(k_out, "int32")
                 else:
                     r = backend.real(backend.cast(status[k], rdtypestr))
                     cdf = backend.cumsum(pu)
-                    k_out = backend.sum(backend.cast(cdf < r, "int32"))
+                    k_out = backend.sum(backend.cast(r >= cdf, "int32"))
                     k_out = backend.clip(
                         k_out,
                         backend.cast(backend.convert_to_tensor(0), "int32"),
diff --git a/tensorcircuit/mpscircuit.py b/tensorcircuit/mpscircuit.py
index c9ec6a9d..b2af7f6b 100644
--- a/tensorcircuit/mpscircuit.py
+++ b/tensorcircuit/mpscircuit.py
@@ -397,7 +397,8 @@ def gate_to_MPO(
         in_dims = tuple(backend.shape_tuple(gate))[:nindex]
         dim = int(in_dims[0])
         dim_phys_mpo = dim * dim
-
+        # transform gate from (in1, in2, ..., out1, out2 ...) to
+        # (in1, out1, in2, out2, ...)
         order = tuple(np.arange(2 * nindex).reshape(2, nindex).T.flatten().tolist())
         gate = backend.transpose(gate, order)
 
@@ -408,14 +409,16 @@ def gate_to_MPO(
             np.column_stack([2 * pair_order, 2 * pair_order + 1])
         ).astype(int)
         pair_axis_perm = tuple(pair_axis_perm.tolist())  # type: ignore
+        # reorder the gate according to the site positions
         gate = backend.transpose(gate, pair_axis_perm)
         index_arr = index_arr[pair_order]  # 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=dim_phys_mpo, norm=False
         )
-
+        # each tensor is in shape of (i, a, b, j)
         tensors: list[Tensor] = []
         previous_i: Optional[int] = None
 
@@ -1042,36 +1045,30 @@ def measure(
         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)
             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))
-
-            cdf = backend.cumsum(ps)
-            choice = backend.sum(backend.cast(r >= cdf, "int32"))
-
-            choice_f = backend.cast(choice, dtypestr)
-            sample.append(choice_f)
-
-            m = backend.zeros((ps.shape[0],), dtype=dtypestr)
-            m = backend.scatter(
-                m,
-                backend.convert_to_tensor([[backend.cast(choice, "int32")]]),
-                backend.convert_to_tensor(np.array([1.0], dtype=dtypestr)),
-            )
-
-            p = p * backend.sum(ps * backend.cast(m, dtype=rdtypestr))
+                r = backend.real(backend.cast(status[k], rdtypestr))
+                cdf = backend.cumsum(ps)
+                eps = 0.31415926 * 1e-12
+                ge_mask = backend.cast(r + eps >= cdf, rdtypestr)
+                outcome = backend.cast(backend.sum(ge_mask), "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:

From e3f88dd11b61e17c3b4abe1193e12108262dd98e Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Tue, 2 Sep 2025 13:40:48 +0800
Subject: [PATCH 42/55] fixed a possible bug

---
 tensorcircuit/basecircuit.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tensorcircuit/basecircuit.py b/tensorcircuit/basecircuit.py
index cda2fdb7..414bf8dd 100644
--- a/tensorcircuit/basecircuit.py
+++ b/tensorcircuit/basecircuit.py
@@ -475,7 +475,9 @@ def measure_jit(
                 pu = pu / backend.sum(pu)
                 if status is None:
                     k_out = backend.implicit_randc(
-                        self._d, shape=1, p=backend.cast(pu, rdtypestr)
+                        a=backend.arange(self._d),
+                        shape=1,
+                        p=backend.cast(pu, rdtypestr),
                     )[0]
                     k_out = backend.cast(k_out, "int32")
                 else:

From 810d624e42ce7e2fccad6a08f6c9c01fff9a4bfb Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Tue, 2 Sep 2025 13:45:28 +0800
Subject: [PATCH 43/55] adjust n_max_d for dim-dimensional qudit circuits.

---
 tensorcircuit/quantum.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tensorcircuit/quantum.py b/tensorcircuit/quantum.py
index 9e324bfc..cc46e594 100644
--- a/tensorcircuit/quantum.py
+++ b/tensorcircuit/quantum.py
@@ -2976,10 +2976,11 @@ def sample2all(
     :rtype: Any
     """
     dim = 2 if dim is None else int(dim)
-    if n > 32:
+    n_max_d = int(32 / np.log2(dim))
+    if n > n_max_d:
         assert (
             len(backend.shape_tuple(sample)) == 2
-        ), "n>32 is only supported for ``sample_bin``"
+        ), f"n>{n_max_d} is only supported for ``sample_bin``"
         if format == "sample_bin":
             return sample
         if format == "count_dict_bin":

From 801fb3ae52776f887625eb9138d0bc06b356f903 Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Tue, 2 Sep 2025 13:57:42 +0800
Subject: [PATCH 44/55] fix argsort and reshape according to comments.

---
 tensorcircuit/mpscircuit.py | 13 ++-----------
 1 file changed, 2 insertions(+), 11 deletions(-)

diff --git a/tensorcircuit/mpscircuit.py b/tensorcircuit/mpscircuit.py
index b2af7f6b..1a98a39d 100644
--- a/tensorcircuit/mpscircuit.py
+++ b/tensorcircuit/mpscircuit.py
@@ -397,22 +397,12 @@ def gate_to_MPO(
         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().tolist())
         gate = backend.transpose(gate, order)
-
-        index_arr = np.array(index, dtype=int) - index_left
-        pair_order = np.argsort(index_arr)
-
-        pair_axis_perm = np.ravel(
-            np.column_stack([2 * pair_order, 2 * pair_order + 1])
-        ).astype(int)
-        pair_axis_perm = tuple(pair_axis_perm.tolist())  # type: ignore
         # reorder the gate according to the site positions
-        gate = backend.transpose(gate, pair_axis_perm)
-        index_arr = index_arr[pair_order]  # 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(
@@ -421,6 +411,7 @@ def gate_to_MPO(
         # each tensor is in shape of (i, a, b, j)
         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:

From 799921797e8a7da321646820b158cf92c67c2344 Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Tue, 2 Sep 2025 15:29:41 +0800
Subject: [PATCH 45/55] use _infer_num_sites() instead of risky calculation.

---
 tensorcircuit/results/counts.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorcircuit/results/counts.py b/tensorcircuit/results/counts.py
index 772d2cee..84d1e8ce 100644
--- a/tensorcircuit/results/counts.py
+++ b/tensorcircuit/results/counts.py
@@ -150,12 +150,12 @@ def vec2count(vec: Tensor, prune: bool = False, dim: Optional[int] = None) -> ct
     :param dim: Dimensionality of the vector, defaults to 2
     :return: {base-d string key: value}, key length n
     """
-    from ..quantum import count_vector2dict
+    from ..quantum import count_vector2dict, _infer_num_sites
 
     dim = 2 if dim is None else dim
     if isinstance(vec, list):
         vec = np.array(vec)
-    n = int(np.log(int(vec.shape[0])) / np.log(dim) + 1e-9)
+    n = _infer_num_sites(int(vec.shape[0]), dim)
     c: ct = count_vector2dict(vec, n, key="bin", dim=dim)  # type: ignore
     if prune:
         c = {k: v for k, v in c.items() if np.abs(v) >= 1e-8}

From d6f0eb409d4b8e22b8cfba10f5d480bf53b96f16 Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Tue, 2 Sep 2025 15:48:19 +0800
Subject: [PATCH 46/55] move _decode_basis_label() from basecircuit.py to
 quantum.py

---
 tensorcircuit/basecircuit.py | 27 ++-------------------------
 tensorcircuit/mpscircuit.py  |  3 +--
 tensorcircuit/quantum.py     | 26 +++++++++++++++++++++++++-
 3 files changed, 28 insertions(+), 28 deletions(-)

diff --git a/tensorcircuit/basecircuit.py b/tensorcircuit/basecircuit.py
index 414bf8dd..f4edc016 100644
--- a/tensorcircuit/basecircuit.py
+++ b/tensorcircuit/basecircuit.py
@@ -25,9 +25,10 @@
     sample_int2bin,
     sample2all,
     _infer_num_sites,
+    _decode_basis_label,
 )
 from .abstractcircuit import AbstractCircuit
-from .cons import npdtype, backend, dtypestr, contractor, rdtypestr, _ALPHABET
+from .cons import npdtype, backend, dtypestr, contractor, rdtypestr
 from .simplify import _split_two_qubit_gate
 from .utils import arg_alias
 
@@ -36,30 +37,6 @@
 Tensor = Any
 
 
-def _decode_basis_label(label: str, n: int, dim: int) -> List[int]:
-    if dim > 36:
-        raise NotImplementedError(
-            f"String basis label supports d<=36 (0–9A–Z). Got dim={dim}. "
-            "Use an integer array/tensor of length n instead."
-        )
-    s = label.upper()
-    if len(s) != n:
-        raise ValueError(f"Basis label length mismatch: expect {n}, got {len(s)}")
-    digits = []
-    for ch in s:
-        if ch not in _ALPHABET:
-            raise ValueError(
-                f"Invalid character '{ch}' in basis label (allowed 0–9A–Z)."
-            )
-        v = _ALPHABET.index(ch)
-        if v >= dim:
-            raise ValueError(
-                f"Digit '{ch}' (= {v}) out of range for base-d with dim={dim}."
-            )
-        digits.append(v)
-    return digits
-
-
 class BaseCircuit(AbstractCircuit):
     _nodes: List[tn.Node]
     _front: List[tn.Edge]
diff --git a/tensorcircuit/mpscircuit.py b/tensorcircuit/mpscircuit.py
index 1a98a39d..b9ab2c6b 100644
--- a/tensorcircuit/mpscircuit.py
+++ b/tensorcircuit/mpscircuit.py
@@ -15,10 +15,9 @@
 
 from . import gates
 from .cons import backend, npdtype, contractor, rdtypestr, dtypestr
-from .quantum import QuOperator, QuVector, extract_tensors_from_qop
+from .quantum import QuOperator, QuVector, extract_tensors_from_qop, _decode_basis_label
 from .mps_base import FiniteMPS
 from .abstractcircuit import AbstractCircuit
-from .basecircuit import _decode_basis_label
 from .utils import arg_alias
 
 Gate = gates.Gate
diff --git a/tensorcircuit/quantum.py b/tensorcircuit/quantum.py
index cc46e594..859580cc 100644
--- a/tensorcircuit/quantum.py
+++ b/tensorcircuit/quantum.py
@@ -32,7 +32,7 @@
     remove_node,
 )
 
-from .cons import backend, contractor, dtypestr, npdtype, rdtypestr
+from .cons import backend, contractor, dtypestr, npdtype, rdtypestr, _ALPHABET
 from .gates import Gate, num_to_tensor
 from .utils import arg_alias
 
@@ -57,6 +57,30 @@ def get_all_nodes(edges: Iterable[Edge]) -> List[Node]:
     return nodes
 
 
+def _decode_basis_label(label: str, n: int, dim: int) -> List[int]:
+    if dim > 36:
+        raise NotImplementedError(
+            f"String basis label supports d<=36 (0–9A–Z). Got dim={dim}. "
+            "Use an integer array/tensor of length n instead."
+        )
+    s = label.upper()
+    if len(s) != n:
+        raise ValueError(f"Basis label length mismatch: expect {n}, got {len(s)}")
+    digits = []
+    for ch in s:
+        if ch not in _ALPHABET:
+            raise ValueError(
+                f"Invalid character '{ch}' in basis label (allowed 0–9A–Z)."
+            )
+        v = _ALPHABET.index(ch)
+        if v >= dim:
+            raise ValueError(
+                f"Digit '{ch}' (= {v}) out of range for base-d with dim={dim}."
+            )
+        digits.append(v)
+    return digits
+
+
 def _infer_num_sites(D: int, dim: int) -> int:
     """
     Infer the number of sites (n) from a Hilbert space dimension D

From ca7867871f926d9a57a5ea1893f6a68397bd8049 Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Tue, 2 Sep 2025 15:52:08 +0800
Subject: [PATCH 47/55] use _infer

---
 tensorcircuit/circuit.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorcircuit/circuit.py b/tensorcircuit/circuit.py
index 0bca86ec..ef982e00 100644
--- a/tensorcircuit/circuit.py
+++ b/tensorcircuit/circuit.py
@@ -87,7 +87,7 @@ def __init__(
             inputs = backend.cast(inputs, dtype=dtypestr)
             inputs = backend.reshape(inputs, [-1])
             N = inputs.shape[0]
-            n = int(np.log(N) / np.log(self._d))
+            n = _infer_num_sites(N, dim=self._d)
             assert n == nqubits or n == 2 * nqubits
             inputs = backend.reshape(inputs, [self._d for _ in range(n)])
             inputs = Gate(inputs)

From bf9e6fc5a4fdfa4fd275faa94f009d982b74ceec Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Tue, 2 Sep 2025 15:53:10 +0800
Subject: [PATCH 48/55] add import

---
 tensorcircuit/circuit.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorcircuit/circuit.py b/tensorcircuit/circuit.py
index ef982e00..18834e01 100644
--- a/tensorcircuit/circuit.py
+++ b/tensorcircuit/circuit.py
@@ -16,7 +16,7 @@
 from . import gates
 from . import channels
 from .cons import backend, contractor, dtypestr, npdtype, _ALPHABET
-from .quantum import QuOperator, identity
+from .quantum import QuOperator, identity, _infer_num_sites
 from .simplify import _full_light_cone_cancel
 from .basecircuit import BaseCircuit
 

From 4e2e27d8ad23eb339690d57384ebd79785c19040 Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Tue, 2 Sep 2025 16:08:33 +0800
Subject: [PATCH 49/55] tiny changes based on comments.

---
 tensorcircuit/densitymatrix.py | 6 ++----
 tensorcircuit/mpscircuit.py    | 2 +-
 tensorcircuit/quantum.py       | 7 ++++++-
 3 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/tensorcircuit/densitymatrix.py b/tensorcircuit/densitymatrix.py
index 690f4e10..5fbfa420 100644
--- a/tensorcircuit/densitymatrix.py
+++ b/tensorcircuit/densitymatrix.py
@@ -222,7 +222,7 @@ def apply_general_kraus(
             dd = dmc.densitymatrix()
             circuits.append(dd)
         tensor = reduce(add, circuits)
-        tensor = backend.reshape(tensor, [self._d for _ in range(2 * self._nqubits)])
+        tensor = backend.reshaped(tensor, d=self._d)
         self._nodes = [Gate(tensor)]
         dangling = [e for e in self._nodes[0]]
         self._front = dangling
@@ -260,9 +260,7 @@ def densitymatrix(self, check: bool = False, reuse: bool = True) -> Tensor:
             t = contractor(nodes, output_edge_order=d_edges)
         else:
             t = nodes[0]
-        dm = backend.reshape(
-            t.tensor, shape=[self._d**self._nqubits, self._d**self._nqubits]
-        )
+        dm = backend.reshapem(t.tensor)
         if check:
             self.check_density_matrix(dm)
         return dm
diff --git a/tensorcircuit/mpscircuit.py b/tensorcircuit/mpscircuit.py
index b9ab2c6b..8011b3bb 100644
--- a/tensorcircuit/mpscircuit.py
+++ b/tensorcircuit/mpscircuit.py
@@ -386,7 +386,7 @@ def gate_to_MPO(
         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 AssertionError("`index` must be strictly increasing.")
+            raise ValueError("`index` must be strictly increasing.")
 
         index_left = int(np.min(index))
         if isinstance(gate, tn.Node):
diff --git a/tensorcircuit/quantum.py b/tensorcircuit/quantum.py
index 859580cc..88cb4d9d 100644
--- a/tensorcircuit/quantum.py
+++ b/tensorcircuit/quantum.py
@@ -2763,7 +2763,7 @@ def sample_bin2int(sample: Tensor, n: int, dim: Optional[int] = None) -> Tensor:
 
     :param sample: in shape [trials, n] of elements (0, 1)
     :type sample: Tensor
-    :param n: number of qubits
+    :param n: number of sites
     :type n: int
     :param dim: local dimension, defaults to 2
     :type dim: int, optional
@@ -2785,10 +2785,15 @@ def sample2count(
     sample_int to count_tuple (indices, counts), size = d**n
 
     :param sample: linear-index samples, shape [shots]
+    :type sample: Tensor
     :param n: number of sites
+    :type n: int
     :param jittable: whether to return fixed-size outputs (backend dependent)
+    :type jittable: bool
     :param dim: local dimension per site, default 2 (qubit)
+    :type dim: int, optional
     :return: (unique_indices, counts)
+    :rtype: Tuple[Tensor, Tensor]
     """
     dim = 2 if dim is None else dim
     size = dim**n

From 793876b46e74970160afbb3ed68a090dfbc0cf59 Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Tue, 2 Sep 2025 16:29:49 +0800
Subject: [PATCH 50/55] use backend.probability_sample()

---
 tensorcircuit/basecircuit.py | 26 ++++++++++++++------------
 tensorcircuit/mpscircuit.py  | 11 ++++++-----
 2 files changed, 20 insertions(+), 17 deletions(-)

diff --git a/tensorcircuit/basecircuit.py b/tensorcircuit/basecircuit.py
index f4edc016..7d14e64e 100644
--- a/tensorcircuit/basecircuit.py
+++ b/tensorcircuit/basecircuit.py
@@ -447,25 +447,27 @@ def measure_jit(
                 sample.append(sign_complex)
                 p = p * (pu * (-1) ** sign + sign)
             else:
-                zero_r = backend.cast(backend.convert_to_tensor(0.0), rdtypestr)
-                pu = backend.clip(backend.real(backend.diagonal(rho)), zero_r, one_r)
+                pu = backend.clip(
+                    backend.real(backend.diagonal(rho)),
+                    backend.convert_to_tensor(0.0),
+                    backend.convert_to_tensor(1.0),
+                )
                 pu = pu / backend.sum(pu)
                 if status is None:
-                    k_out = backend.implicit_randc(
+                    ind = backend.implicit_randc(
                         a=backend.arange(self._d),
                         shape=1,
                         p=backend.cast(pu, rdtypestr),
-                    )[0]
-                    k_out = backend.cast(k_out, "int32")
+                    )
                 else:
-                    r = backend.real(backend.cast(status[k], rdtypestr))
-                    cdf = backend.cumsum(pu)
-                    k_out = backend.sum(backend.cast(r >= cdf, "int32"))
-                    k_out = backend.clip(
-                        k_out,
-                        backend.cast(backend.convert_to_tensor(0), "int32"),
-                        backend.cast(backend.convert_to_tensor(self._d - 1), "int32"),
+                    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(pu, rdtypestr),
+                        status=one_r - st,
                     )
+                k_out = backend.cast(ind[0], "int32")
                 sample.append(backend.cast(k_out, rdtypestr))
                 p = p * backend.cast(pu[k_out], rdtypestr)
         sample = backend.real(backend.stack(sample))
diff --git a/tensorcircuit/mpscircuit.py b/tensorcircuit/mpscircuit.py
index 8011b3bb..7b6e8d7b 100644
--- a/tensorcircuit/mpscircuit.py
+++ b/tensorcircuit/mpscircuit.py
@@ -1048,11 +1048,12 @@ def measure(
                     self._d, shape=1, p=backend.cast(ps, rdtypestr)
                 )[0]
             else:
-                r = backend.real(backend.cast(status[k], rdtypestr))
-                cdf = backend.cumsum(ps)
-                eps = 0.31415926 * 1e-12
-                ge_mask = backend.cast(r + eps >= cdf, rdtypestr)
-                outcome = backend.cast(backend.sum(ge_mask), "int32")
+                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))

From 1c5e67e2a1a7124acd6dc2424e3c823e05e86952 Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Tue, 2 Sep 2025 17:06:08 +0800
Subject: [PATCH 51/55] add onehot_d_tensor in quantum.py and apply.

---
 tensorcircuit/basecircuit.py | 25 ++++++------------------
 tensorcircuit/quantum.py     | 37 ++++++++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 19 deletions(-)

diff --git a/tensorcircuit/basecircuit.py b/tensorcircuit/basecircuit.py
index 7d14e64e..f996cf08 100644
--- a/tensorcircuit/basecircuit.py
+++ b/tensorcircuit/basecircuit.py
@@ -26,6 +26,7 @@
     sample2all,
     _infer_num_sites,
     _decode_basis_label,
+    onehot_d_tensor,
 )
 from .abstractcircuit import AbstractCircuit
 from .cons import npdtype, backend, dtypestr, contractor, rdtypestr
@@ -410,8 +411,7 @@ def measure_jit(
                         np.array([1, 0])
                     ) + sample[i] * gates.array_to_tensor(np.array([0, 1]))
                 else:
-                    vec = backend.one_hot(backend.cast(sample[i], "int32"), self._d)
-                    m = backend.cast(vec, dtypestr)
+                    m = onehot_d_tensor(sample[i], d=self._d)
                 g1 = Gate(m)
                 g1.id = id(g1)
                 g1.is_dagger = False
@@ -507,11 +507,6 @@ def amplitude_before(self, l: Union[str, Tensor]) -> List[Gate]:
         :rtype: List[Gate]
         """
 
-        def _basis_nod(_k: int) -> Tensor:
-            _vec = np.zeros((self._d,), dtype=npdtype)
-            _vec[_k] = 1.0
-            return _vec
-
         no, d_edges = self._copy()
         ms = []
         if self.is_dm:
@@ -519,17 +514,14 @@ def _basis_nod(_k: int) -> Tensor:
         if isinstance(l, str):
             symbols = _decode_basis_label(l, n=self._nqubits, dim=self._d)
             for k in symbols:
-                n = _basis_nod(k)
+                n = onehot_d_tensor(k, d=self._d)
                 ms.append(tn.Node(n))
                 if self.is_dm:
                     msconj.append(tn.Node(n))
         else:
             l = backend.cast(l, dtype=dtypestr)
             for i in range(self._nqubits):
-                endn = backend.cast(
-                    backend.one_hot(backend.cast(l[i], "int32"), self._d),
-                    dtype=dtypestr,
-                )
+                endn = onehot_d_tensor(l[i], d=self._d)
                 ms.append(tn.Node(endn))
                 if self.is_dm:
                     msconj.append(tn.Node(endn))
@@ -1040,18 +1032,13 @@ def projected_subsystem(self, traceout: Tensor, left: Tuple[int, ...]) -> Tensor
         :rtype: Tensor
         """
 
-        def _basis_gate(k_tensor: Any) -> Gate:
-            vec = backend.one_hot(backend.cast(k_tensor, "int32"), self._d)
-            vec = backend.cast(vec, dtypestr)
-            return Gate(vec)
-
         traceout = backend.cast(traceout, dtypestr)
         nodes, front = self._copy()
         L = self._nqubits
         edges = []
         for i in range(len(traceout)):
             if i not in left:
-                n = _basis_gate(traceout[i])
+                n = Gate(onehot_d_tensor(traceout[i], d=self._d))
                 nodes.append(n)
                 front[i] ^ n[0]
             else:
@@ -1060,7 +1047,7 @@ def _basis_gate(k_tensor: Any) -> Gate:
         if self.is_dm:
             for i in range(len(traceout)):
                 if i not in left:
-                    n = _basis_gate(traceout[i])
+                    n = Gate(onehot_d_tensor(traceout[i], d=self._d))
                     nodes.append(n)
                     front[i + L] ^ n[0]
                 else:
diff --git a/tensorcircuit/quantum.py b/tensorcircuit/quantum.py
index 88cb4d9d..1755a7bb 100644
--- a/tensorcircuit/quantum.py
+++ b/tensorcircuit/quantum.py
@@ -57,7 +57,44 @@ def get_all_nodes(edges: Iterable[Edge]) -> List[Node]:
     return nodes
 
 
+def onehot_d_tensor(_k: Union[int, Tensor], d: int = 2) -> Tensor:
+    """
+    Construct a one-hot vector (or matrix) of local dimension ``d``.
+
+    :param _k: index or indices to set as 1. Can be an int or a backend Tensor.
+    :type _k: int or Tensor
+    :param d: local dimension (number of categories), defaults to 2
+    :type d: int, optional
+    :return: one-hot encoded vector (shape [d]) or matrix (shape [len(_k), d])
+    :rtype: Tensor
+    """
+    if isinstance(_k, int):
+        vec = backend.one_hot(_k, d)
+    else:
+        vec = backend.one_hot(backend.cast(_k, "int32"), d)
+    return backend.cast(vec, dtypestr)
+
+
 def _decode_basis_label(label: str, n: int, dim: int) -> List[int]:
+    """
+    Decode a string basis label into a list of integer digits.
+
+    The label is interpreted in base-``dim`` using characters ``0–9A–Z``.
+    Only dimensions up to 36 are supported.
+
+    :param label: basis label string, e.g. "010" or "A9F"
+    :type label: str
+    :param n: number of sites (expected length of the label)
+    :type n: int
+    :param dim: local dimension (2 <= dim <= 36)
+    :type dim: int
+    :return: list of integer digits of length ``n``, each in ``[0, dim-1]``
+    :rtype: List[int]
+
+    :raises NotImplementedError: if ``dim > 36``
+    :raises ValueError: if the label length mismatches ``n``,
+                        or contains invalid/out-of-range characters
+    """
     if dim > 36:
         raise NotImplementedError(
             f"String basis label supports d<=36 (0–9A–Z). Got dim={dim}. "

From 78120ad2089819359303fd1c1d7c2f23f9af0c38 Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Tue, 2 Sep 2025 17:24:28 +0800
Subject: [PATCH 52/55] remove parse_key

---
 tensorcircuit/results/counts.py | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/tensorcircuit/results/counts.py b/tensorcircuit/results/counts.py
index 84d1e8ce..990361c2 100644
--- a/tensorcircuit/results/counts.py
+++ b/tensorcircuit/results/counts.py
@@ -6,7 +6,7 @@
 
 import numpy as np
 
-from ..cons import _ALPHABET
+from ..quantum import _decode_basis_label
 
 Tensor = Any
 ct = Dict[str, int]
@@ -117,9 +117,6 @@ def count2vec(
     array([0.2, 0. , 0.3, 0.5])
     """
 
-    def parse_key(_k: str) -> List[int]:
-        return [_ALPHABET.index(_ch) for _ch in _k.upper()]
-
     if not count:
         return np.array([], dtype=float)
 
@@ -133,7 +130,7 @@ def parse_key(_k: str) -> List[int]:
 
     powers = [dim**p for p in range(n)][::-1]
     for k, v in count.items():
-        digits = parse_key(k)
+        digits = _decode_basis_label(k, n, dim)
         idx = sum(dig * p for dig, p in zip(digits, powers))
         prob[idx] = (v / shots) if normalization else v
 

From d84bd7265cfa7418dedd53ccc6f6de1f5de5b030 Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Tue, 2 Sep 2025 17:28:33 +0800
Subject: [PATCH 53/55] use onehot_d_tensor

---
 tensorcircuit/circuit.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/tensorcircuit/circuit.py b/tensorcircuit/circuit.py
index 18834e01..4c1312ff 100644
--- a/tensorcircuit/circuit.py
+++ b/tensorcircuit/circuit.py
@@ -798,9 +798,7 @@ def measure_reference(
                 if i != j:
                     e ^ edge2[i]
             for i in range(len(sample)):
-                m = np.array([0 for _ in range(self._d)], dtype=npdtype)
-                m[int(sample[i])] = 1
-
+                m = onehot_d_tensor(sample[i], d=self._d)
                 nodes1.append(tn.Node(m))
                 nodes1[-1].get_edge(0) ^ edge1[index[i]]
                 nodes2.append(tn.Node(m))

From a26be1fabb3a78c59468157af735d55a121ea685 Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Tue, 2 Sep 2025 17:28:54 +0800
Subject: [PATCH 54/55] import onhot_d_tensor

---
 tensorcircuit/circuit.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorcircuit/circuit.py b/tensorcircuit/circuit.py
index 4c1312ff..2347bdca 100644
--- a/tensorcircuit/circuit.py
+++ b/tensorcircuit/circuit.py
@@ -16,7 +16,7 @@
 from . import gates
 from . import channels
 from .cons import backend, contractor, dtypestr, npdtype, _ALPHABET
-from .quantum import QuOperator, identity, _infer_num_sites
+from .quantum import QuOperator, identity, _infer_num_sites, onehot_d_tensor
 from .simplify import _full_light_cone_cancel
 from .basecircuit import BaseCircuit
 

From f82b382168f021edb2dcb0df449254eb656bbeac Mon Sep 17 00:00:00 2001
From: Weiguo Ma <david.mawg@gmail.com>
Date: Tue, 2 Sep 2025 17:35:58 +0800
Subject: [PATCH 55/55] removed a useless import

---
 tensorcircuit/results/counts.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorcircuit/results/counts.py b/tensorcircuit/results/counts.py
index 990361c2..662da8d9 100644
--- a/tensorcircuit/results/counts.py
+++ b/tensorcircuit/results/counts.py
@@ -2,7 +2,7 @@
 dict related functionalities
 """
 
-from typing import Any, Dict, Optional, Sequence, List
+from typing import Any, Dict, Optional, Sequence
 
 import numpy as np