tensorcircuit-nightly 1.3.0.dev20250903__py3-none-any.whl → 1.3.0.dev20250905__py3-none-any.whl

tensorcircuit/quantum.py

@@ -9,6 +9,7 @@ import math
 import os
 from functools import partial, reduce
 from operator import matmul, mul, or_
+from collections import Counter
 from typing import (
     Any,
     Callable,
@@ -31,7 +32,7 @@ from tensornetwork.network_operations import (
     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
 
@@ -56,6 +57,91 @@ 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}. "
+            "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
+    and local dimension d, assuming D = d**n.
+
+    :param D: total Hilbert space dimension (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(dim, int) and dim >= 2):
+        raise ValueError(f"d must be an integer >= 2, got {dim}")
+
+    tmp, n = D, 0
+    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 {dim}")
+    return n
+
+
 def _reachable(nodes: List[AbstractNode]) -> List[AbstractNode]:
     if not nodes:
         raise ValueError("Reachable requires at least 1 node.")
@@ -2150,7 +2236,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,
+    dim: Optional[int] = None,
 ) -> Tensor:
     """
     Compute the reduced wavefunction from the quantum state ``state``.
@@ -2165,20 +2254,22 @@ def reduced_wavefunction(
     :type measure: List[int]
     :return: _description_
     :rtype: Tensor
+    :param dim: dimension of qudit system
+    :type dim: int
     """
+    dim = 2 if dim is None else dim
     if measure is None:
         measure = [0 for _ in cut]
-    s = backend.reshape2(state)
+    s = backend.reshaped(state, dim)
     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"), dim),
+            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 +2284,9 @@ def reduced_density_matrix(
     cut: Union[int, List[int]],
     p: Optional[Tensor] = None,
     normalize: bool = True,
+    dim: 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 +2298,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 dim: dimension of qudit system
+    :type dim: int
     """
+    dim = 2 if dim is None else dim
     if isinstance(cut, list) or isinstance(cut, tuple) or isinstance(cut, set):
         traceout = list(cut)
     else:
@@ -2220,21 +2316,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], 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, [2 for _ in range(2 * freedom)])
+
+        rho = backend.reshape(state, [dim] * (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)),
+                dim ** len(traceout),
+                dim ** (freedom - len(traceout)),
+                dim ** len(traceout),
+                dim ** (freedom - len(traceout)),
             ],
         )
         if p is None:
@@ -2247,20 +2341,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, [dim ** (freedom - len(traceout)), dim ** (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, dim)
         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, [dim for _ in range(freedom)])
         w = backend.transpose(w, perm=perm)
-        w = backend.reshape(w, [-1, 2 ** len(traceout)])
+        w = backend.reshape(w, [-1, dim ** len(traceout)])
         if p is None:
             rho = w @ backend.adjoint(w)
         else:
@@ -2403,7 +2497,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], dim: Optional[int] = None
+) -> Tensor:
     """
     _summary_
 
@@ -2411,10 +2507,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 dim: dimension of qudit system
+    :type dim: int
     :return: _description_
     :rtype: Tensor
     """
-    rho = backend.reshape2(rho)
+    dim = 2 if dim is None else dim
+    rho = backend.reshaped(rho, dim)
     rho_node = Gate(rho)
     n = len(rho.shape) // 2
     left_edges = []
@@ -2432,7 +2531,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], dim: Optional[int] = None
+) -> Tensor:
     """
     _summary_
 
@@ -2440,17 +2541,21 @@ def entanglement_negativity(rho: Tensor, transposed_sites: List[int]) -> Tensor:
     :type rho: Tensor
     :param transposed_sites: _description_
     :type transposed_sites: List[int]
+    :param dim: dimension of qudit system
+    :type dim: int
     :return: _description_
     :rtype: Tensor
     """
-    rhot = partial_transpose(rho, transposed_sites)
+    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
 
 
 @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", dim: Optional[int] = None
+) -> Tensor:
     """
     _summary_
 
@@ -2460,10 +2565,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 dim: dimension of qudit system
+    :type dim: int
     :return: _description_
     :rtype: Tensor
     """
-    rhot = partial_transpose(rho, transposed_sites)
+    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)
@@ -2549,7 +2657,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]], dim: Optional[int] = None
+) -> Tensor:
     """
     Mutual information between AB subsystem described by ``cut``.
 
@@ -2557,9 +2667,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 dim: The diagonal matrix in form of Tensor.
+    :type dim: Tensor
     :return: The mutual information between AB subsystem described by ``cut``.
     :rtype: Tensor
     """
+    dim = 2 if dim is None else dim
     if isinstance(cut, list) or isinstance(cut, tuple) or isinstance(cut, set):
         traceout = list(cut)
     else:
@@ -2567,22 +2680,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], dim=dim)
         hab = entropy(s)
 
         # subsystem a
-        rhoa = reduced_density_matrix(s, traceout)
+        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)  # 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)
+        rhoa = reduced_density_matrix(s, traceout, dim=dim)
         ha = hb = entropy(rhoa)
 
     return ha + hb - hab
@@ -2591,7 +2704,9 @@ 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, dim: Optional[int] = None
+) -> Tensor:
     """
     measurement shots results, sparse tuple representation to dense representation
     count_vector to count_tuple
@@ -2600,11 +2715,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 dim: [description], defaults to None
+    :type dim: int, optional
     :return: [description]
     :rtype: Tensor
     """
+    dim = 2 if dim is None else dim
     return backend.scatter(
-        backend.cast(backend.zeros([2**n]), srepr[1].dtype),
+        backend.cast(backend.zeros([dim**n]), srepr[1].dtype),
         backend.reshape(srepr[0], [-1, 1]),
         srepr[1],
     )
@@ -2647,117 +2765,146 @@ 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, dim: 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 dim: local dimension, defaults to 2
+    :type dim: 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
+    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(dim, pos)
+        digits = backend.mod(
+            backend.floor_divide(sample[..., None], base),  # ⌊sample / d**pos⌋
+            dim,
+        )
+        return backend.cast(digits, "int32")
 
 
-def sample_bin2int(sample: Tensor, n: int) -> Tensor:
+def sample_bin2int(sample: Tensor, n: int, dim: Optional[int] = None) -> Tensor:
     """
     bin sample to int sample
 
     :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
     :return: in shape [trials]
     :rtype: Tensor
     """
-    power = backend.convert_to_tensor([2**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)
 
 
 def sample2count(
-    sample: Tensor, n: int, jittable: bool = True
+    sample: Tensor,
+    n: int,
+    jittable: bool = True,
+    dim: Optional[int] = None,
 ) -> Tuple[Tensor, Tensor]:
     """
-    sample_int to count_tuple
+    sample_int to count_tuple (indices, counts), size = d**n
 
-    :param sample: _description_
+    :param sample: linear-index samples, shape [shots]
     :type sample: Tensor
-    :param n: _description_
+    :param n: number of sites
     :type n: int
-    :param jittable: _description_, defaults to True
-    :type jittable: bool, optional
-    :return: _description_
+    :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]
     """
-    d = 2**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
-        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", dim: 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 dim: local dimension (default 2)
+    :type dim: int, optional
+    :return: mapping from configuration to count
+    :rtype: Dict[Any, int]
     """
     from .interfaces import which_backend
 
+    dim = 2 if dim is None else dim
     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(dim**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=dim).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", dim: 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 dim: local dimension, defaults to 2
+    :type dim: int, optional
     :return: count_dict
-    :rtype: _type_
+    :rtype: Dict[Any, int]
     """
-    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])
         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=dim).zfill(n)
+            out_str[kn] = v
+        return out_str
 
 
 @partial(arg_alias, alias_dict={"counts": ["shots"], "format": ["format_"]})
@@ -2769,8 +2916,9 @@ def measurement_counts(
     random_generator: Optional[Any] = None,
     status: Optional[Tensor] = None,
     jittable: bool = False,
+    dim: 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 +2933,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 +2985,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 dim is None else dim
+    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", dim=local_d)
         elif format == "count_dict_int":
-            return count_vector2dict(pi, n, key="int")
+            return count_vector2dict(pi, n, key="int", dim=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 +3011,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, dim=local_d)
 
 
 measurement_results = measurement_counts
@@ -2869,45 +3019,64 @@ measurement_results = 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,
+    dim: 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 dim: local dimension (2 for qubit; >2 for qudit), defaults to 2
+    :type dim: Optional[int]
     :return: measurement results specified as ``format``
     :rtype: Any
     """
+    dim = 2 if dim is None else int(dim)
+    n_max_d = int(32 / np.log2(dim))
+    if n > n_max_d:
+        assert (
+            len(backend.shape_tuple(sample)) == 2
+        ), f"n>{n_max_d} 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
-        sample_bin = sample_int2bin(sample, n)
+        sample_bin = sample_int2bin(sample, n, dim=dim)
     elif len(backend.shape_tuple(sample)) == 2:
-        sample_int = sample_bin2int(sample, n)
+        sample_int = sample_bin2int(sample, n, dim=dim)
         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, dim=dim)
         if format == "count_tuple":
             return count_tuple
         elif format == "count_vector":
-            return count_s2d(count_tuple, n)
+            return count_s2d(count_tuple, n, dim=dim)
         elif format == "count_dict_bin":
-            return count_tuple2dict(count_tuple, n, key="bin")
+            return count_tuple2dict(count_tuple, n, key="bin", dim=dim)
         elif format == "count_dict_int":
-            return count_tuple2dict(count_tuple, n, key="int")
+            return count_tuple2dict(count_tuple, n, key="int", dim=dim)
         else:
             raise ValueError(
                 "unsupported format %s for finite shots measurement" % format