tensorcircuit-nightly 1.0.2.dev20250108__py3-none-any.whl → 1.4.0.dev20251103__py3-none-any.whl

tensorcircuit/circuit.py

@@ -1,5 +1,7 @@
 """
-Quantum circuit: the state simulator
+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.
 """
 
 # pylint: disable=invalid-name
@@ -13,8 +15,8 @@ import tensornetwork as tn
 
 from . import gates
 from . import channels
-from .cons import backend, contractor, dtypestr, npdtype
-from .quantum import QuOperator, identity
+from .cons import backend, contractor, dtypestr, npdtype, _ALPHABET
+from .quantum import QuOperator, identity, _infer_num_sites, onehot_d_tensor
 from .simplify import _full_light_cone_cancel
 from .basecircuit import BaseCircuit
 
@@ -23,7 +25,7 @@ Tensor = Any
 
 
 class Circuit(BaseCircuit):
-    """
+    r"""
     ``Circuit`` class.
     Simple usage demo below.
 
@@ -45,14 +47,18 @@ class Circuit(BaseCircuit):
         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.
+        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
+        :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,7 @@ class Circuit(BaseCircuit):
             ``max_singular_values`` and ``max_truncation_err``.
         :type split: Optional[Dict[str, Any]]
         """
+        self._d = 2 if dim is None else dim
         self.inputs = inputs
         self.mps_inputs = mps_inputs
         self.split = split
@@ -70,18 +77,19 @@ class Circuit(BaseCircuit):
             "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)
+            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)
             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, dim=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 +186,14 @@ class Circuit(BaseCircuit):
 
         :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)
@@ -231,6 +226,25 @@ class Circuit(BaseCircuit):
         pz: float,
         status: Optional[float] = None,
     ) -> float:
+        """
+        Apply a depolarizing channel to the circuit in a Monte Carlo way.
+        For each call, one of the Pauli gates (X, Y, Z) or an Identity gate is applied to the qubit
+        at the given index based on the probabilities `px`, `py`, and `pz`.
+
+        :param index: The index of the qubit to apply the depolarizing channel on.
+        :type index: int
+        :param px: The probability of applying an X gate.
+        :type px: float
+        :param py: The probability of applying a Y gate.
+        :type py: float
+        :param pz: The probability of applying a Z gate.
+        :type pz: float
+        :param status: A random number between 0 and 1 to determine which gate to apply. If None,
+                       a random number is generated automatically. Defaults to None.
+        :type status: Optional[float], optional
+        :return: Returns 0.0. The function modifies the circuit in place.
+        :rtype: float
+        """
         if status is None:
             status = backend.implicit_randu()[0]
         g = backend.cond(
@@ -323,6 +337,35 @@ class Circuit(BaseCircuit):
         status: Optional[float] = None,
         name: Optional[str] = None,
     ) -> Tensor:
+        """
+        Apply a unitary Kraus channel to the circuit using a Monte Carlo approach. This method is functionally
+        similar to `unitary_kraus` but uses `backend.switch` for selecting the Kraus operator, which can have
+        different performance characteristics on some backends.
+
+        A random Kraus operator from the provided list is applied to the circuit based on the given probabilities.
+        This method is jittable and suitable for simulating noisy quantum circuits where the noise is represented
+        by unitary Kraus operators.
+
+        .. warning::
+            This method may have issues with `vmap` due to potential concurrent access locks, potentially related with
+            `backend.switch`. `unitary_kraus` is generally recommended.
+
+        :param kraus: A sequence of `Gate` objects representing the unitary Kraus operators.
+        :type kraus: Sequence[Gate]
+        :param index: The qubit indices on which to apply the Kraus channel.
+        :type index: int
+        :param prob: A sequence of probabilities corresponding to each Kraus operator. If None, probabilities
+                     are derived from the operators themselves. Defaults to None.
+        :type prob: Optional[Sequence[float]], optional
+        :param status: A random number between 0 and 1 to determine which Kraus operator to apply. If None,
+                       a random number is generated automatically. Defaults to None.
+        :type status: Optional[float], optional
+        :param name: An optional name for the operation. Defaults to None.
+        :type name: Optional[str], optional
+        :return: A tensor indicating which Kraus operator was applied.
+        :rtype: Tensor
+        """
+
         # dont use, has issue conflicting with vmap, concurrent access lock emerged
         # potential issue raised from switch
         # general impl from Monte Carlo trajectory depolarizing above
@@ -431,8 +474,8 @@ class Circuit(BaseCircuit):
         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)])
-        self.any(*index, unitary=g, name=name)  # type: ignore
+        g = backend.reshape(g, [self._d for _ in range(sites * 2)])
+        self.any(*index, unitary=g, name=name, dim=self._d)  # type: ignore
         return r
 
     def _general_kraus_tf(
@@ -557,9 +600,13 @@ class Circuit(BaseCircuit):
             for w, k in zip(prob, kraus_tensor)
         ]
         pick = self.unitary_kraus(
-            new_kraus, *index, prob=prob, status=status, name=name
+            new_kraus,
+            *index,
+            prob=prob,
+            status=status,
+            name=name,
         )
-        if with_prob is False:
+        if not with_prob:
             return pick
         else:
             return pick, prob
@@ -590,7 +637,11 @@ class Circuit(BaseCircuit):
         :type status: Optional[float], optional
         """
         return self._general_kraus_2(
-            kraus, *index, status=status, with_prob=with_prob, name=name
+            kraus,
+            *index,
+            status=status,
+            with_prob=with_prob,
+            name=name,
         )
 
     apply_general_kraus = general_kraus
@@ -632,7 +683,7 @@ class Circuit(BaseCircuit):
             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
@@ -689,8 +740,8 @@ class Circuit(BaseCircuit):
         :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, dim=self._d)
         ns, es = self._copy()
         c._nodes = ns
         c._front = es
@@ -710,8 +761,8 @@ class Circuit(BaseCircuit):
         :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, dim=self._d)
         ns, es = self._copy()
         c._nodes = ns
         c._front = es
@@ -724,6 +775,9 @@ class Circuit(BaseCircuit):
         """
         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)
@@ -752,10 +806,7 @@ class Circuit(BaseCircuit):
                 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 = 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))
@@ -766,15 +817,13 @@ class Circuit(BaseCircuit):
                 / 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 = backend.implicit_randc(self._d, shape=1, p=probs)
+
+            sample += _ALPHABET[outcome]
+            p *= float(probs[outcome])
+
         if with_prob:
             return sample, p
         else:
@@ -794,6 +843,10 @@ class Circuit(BaseCircuit):
     ) -> 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:
 
@@ -829,14 +882,12 @@ class Circuit(BaseCircuit):
         :param nmc: repetition time for Monte Carlo sampling for noisfy calculation, defaults to 1000
         :type nmc: int, optional
         :param status: external randomness given by tensor uniformly from [0, 1], defaults to None,
-            used for noisfy circuit sampling
+            used for noisy circuit sampling
         :type status: Optional[Tensor], optional
         :raises ValueError: "Cannot measure two operators in one index"
         :return: Tensor with one element
         :rtype: Tensor
         """
-        from .noisemodel import expectation_noisfy
-
         if noise_conf is None:
             # if not reuse:
             #     nodes1, edge1 = self._copy()
@@ -851,6 +902,8 @@ class Circuit(BaseCircuit):
                 nodes1 = _full_light_cone_cancel(nodes1)
             return contractor(nodes1).tensor
         else:
+            from .noisemodel import expectation_noisfy
+
             return expectation_noisfy(
                 self,
                 *ops,
@@ -871,9 +924,11 @@ def expectation(
     bra: Optional[Tensor] = None,
     conj: bool = True,
     normalization: bool = False,
+    dim: Optional[int] = None,
 ) -> 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`)
 
@@ -918,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 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.
@@ -926,6 +983,7 @@ def expectation(
     :return: The result of :math:`\\langle bra\\vert ops \\vert ket\\rangle`.
     :rtype: Tensor
     """
+    dim = 2 if dim is None else dim
     if bra is None:
         bra = ket
     if isinstance(ket, QuOperator):
@@ -939,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, dim)
                 op = gates.Gate(op)
             if isinstance(index, int):
                 index = [index]
@@ -963,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, dim)
+        bra = backend.reshaped(bra, dim)
         n = len(backend.shape_tuple(ket))
         ket = Gate(ket)
         bra = Gate(bra)
@@ -976,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, dim)
                 op = gates.Gate(op)
             if isinstance(index, int):
                 index = [index]

tensorcircuit/cloud/local.py

@@ -36,7 +36,7 @@ def submit_task(
     shots: Union[int, Sequence[int]] = 1024,
     version: str = "1",
     circuit: Optional[Union[AbstractCircuit, Sequence[AbstractCircuit]]] = None,
-    **kws: Any
+    **kws: Any,
 ) -> List[Task]:
     def _circuit2result(c: AbstractCircuit) -> Dict[str, Any]:
         if device.name in ["testing", "default"]:

tensorcircuit/cloud/quafu_provider.py

@@ -30,7 +30,7 @@ def submit_task(
     circuit: Optional[Union[AbstractCircuit, Sequence[AbstractCircuit]]] = None,
     source: Optional[Union[str, Sequence[str]]] = None,
     compile: bool = True,
-    **kws: Any
+    **kws: Any,
 ) -> Task:
     if source is None:
 

tensorcircuit/cloud/tencent.py

@@ -133,7 +133,7 @@ def submit_task(
     enable_qos_gate_decomposition: bool = True,
     enable_qos_initial_mapping: bool = False,
     qos_dry_run: bool = False,
-    **kws: Any
+    **kws: Any,
 ) -> List[Task]:
     """
     Submit task via tencent provider, we suggest to enable one of the compiling functionality:

tensorcircuit/compiler/simple_compiler.py

@@ -109,7 +109,7 @@ def prune(
     circuit: Union[AbstractCircuit, List[Dict[str, Any]]],
     rtol: float = 1e-3,
     atol: float = 1e-3,
-    **kws: Any
+    **kws: Any,
 ) -> Any:
     if isinstance(circuit, list):
         qir = circuit
@@ -251,7 +251,7 @@ def _merge(
 def merge(
     circuit: Union[AbstractCircuit, List[Dict[str, Any]]],
     rules: Optional[Dict[Tuple[str, ...], str]] = None,
-    **kws: Any
+    **kws: Any,
 ) -> Any:
     merge_rules = copy(default_merge_rules)
     if rules is not None:

tensorcircuit/cons.py

@@ -8,7 +8,7 @@ import logging
 import sys
 import time
 from contextlib import contextmanager
-from functools import partial, reduce, wraps
+from functools import partial, reduce, wraps, lru_cache
 from operator import mul
 from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence, Tuple
 
@@ -23,6 +23,39 @@ from .simplify import _multi_remove
 
 logger = logging.getLogger(__name__)
 
+## monkey patch
+_NODE_CREATION_COUNTER = 0
+_original_node_init = tn.Node.__init__
+
+
+@wraps(_original_node_init)
+def _patched_node_init(self: Any, *args: Any, **kwargs: Any) -> None:
+    """Patched Node.__init__ to add a stable creation ID."""
+    global _NODE_CREATION_COUNTER
+    _original_node_init(self, *args, **kwargs)
+    self._stable_id_ = _NODE_CREATION_COUNTER
+    _NODE_CREATION_COUNTER += 1
+
+
+tn.Node.__init__ = _patched_node_init
+
+
+def _get_edge_stable_key(edge: tn.Edge) -> Tuple[int, int, int, int]:
+    n1, n2 = edge.node1, edge.node2
+    id1 = getattr(n1, "_stable_id_", -1)
+    id2 = getattr(n2, "_stable_id_", -1) if n2 is not None else -2  # -2 for dangling
+
+    if id1 > id2 or (id1 == id2 and edge.axis1 > edge.axis2):
+        id1, id2, ax1, ax2 = id2, id1, edge.axis2, edge.axis1
+    else:
+        ax1, ax2 = edge.axis1, edge.axis2
+    return (id1, ax1, id2, ax2)
+
+
+def sorted_edges(edges: Iterator[tn.Edge]) -> List[tn.Edge]:
+    return sorted(edges, key=_get_edge_stable_key)
+
+
 package_name = "tensorcircuit"
 thismodule = sys.modules[__name__]
 dtypestr = "complex64"
@@ -30,6 +63,7 @@ rdtypestr = "float32"
 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
 
 
@@ -261,7 +295,7 @@ def _merge_single_gates(
             e0 = n0[0]
         njs = [i for i, n in enumerate(nodes) if id(n) in [id(e0.node1), id(e0.node2)]]
         qjs = [i for i, n in enumerate(queue) if id(n) in [id(e0.node1), id(e0.node2)]]
-        new_node = tn.contract(e0)
+        new_node = tn.contract_parallel(e0)
         total_size += _sizen(new_node)
 
         logger.debug(
@@ -439,6 +473,28 @@ def tn_greedy_contractor(
 # base = tn.contractors.opt_einsum_paths.path_contractors.base
 # utils = tn.contractors.opt_einsum_paths.utils
 
+_einsum_symbols_base = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
+
+
+@lru_cache(2**14)
+def get_symbol(i: int) -> str:
+    """Get the symbol corresponding to int ``i`` - runs through the usual 52
+    letters before resorting to unicode characters, starting at ``chr(192)``
+    and skipping surrogates. From cotengra codebase
+    """
+    if i < 52:
+        # use a-z, A-Z first
+        return _einsum_symbols_base[i]
+
+    # then proceed from 'À'
+    i += 140
+
+    if i >= 55296:
+        # Skip chr(57343) - chr(55296) as surrogates
+        i += 2048
+
+    return chr(i)
+
 
 def _get_path(
     nodes: List[tn.Node], algorithm: Any
@@ -451,30 +507,33 @@ def _get_path(
     return algorithm(input_sets, output_set, size_dict), nodes
 
 
+def _identity(*args: Any, **kws: Any) -> Any:
+    return args
+
+
 def _get_path_cache_friendly(
     nodes: List[tn.Node], algorithm: Any
 ) -> Tuple[List[Tuple[int, int]], List[tn.Node]]:
     nodes = list(nodes)
+
+    nodes_new = sorted(nodes, key=lambda node: getattr(node, "_stable_id_", -1))
+    # if isinstance(algorithm, list):
+    #     return algorithm, [nodes_new]
+
+    all_edges = tn.get_all_edges(nodes_new)
+    all_edges_sorted = sorted_edges(all_edges)
     mapping_dict = {}
     i = 0
-    for n in nodes:
-        for e in n:
-            if id(e) not in mapping_dict:
-                mapping_dict[id(e)] = i
-                i += 1
-    # TODO(@refraction-ray): may be not that cache friendly, since the edge id correspondence is not that fixed?
-    input_sets = [set([mapping_dict[id(e)] for e in node.edges]) for node in nodes]
-    placeholder = [[1e20 for _ in range(100)]]
-    order = np.argsort(np.array(list(map(sorted, input_sets)) + placeholder, dtype=object))[:-1]  # type: ignore
-    nodes_new = [nodes[i] for i in order]
-    if isinstance(algorithm, list):
-        return algorithm, nodes_new
-
-    input_sets = [set([mapping_dict[id(e)] for e in node.edges]) for node in nodes_new]
-    output_set = set([mapping_dict[id(e)] for e in tn.get_subgraph_dangling(nodes_new)])
-    size_dict = {
-        mapping_dict[id(edge)]: edge.dimension for edge in tn.get_all_edges(nodes_new)
-    }
+    for edge in all_edges_sorted:
+        if id(edge) not in mapping_dict:
+            mapping_dict[id(edge)] = get_symbol(i)
+            i += 1
+
+    input_sets = [list([mapping_dict[id(e)] for e in node.edges]) for node in nodes_new]
+    output_set = list(
+        [mapping_dict[id(e)] for e in sorted_edges(tn.get_subgraph_dangling(nodes_new))]
+    )
+    size_dict = {mapping_dict[id(edge)]: edge.dimension for edge in all_edges_sorted}
     logger.debug("input_sets: %s" % input_sets)
     logger.debug("output_set: %s" % output_set)
     logger.debug("size_dict: %s" % size_dict)
@@ -483,6 +542,9 @@ def _get_path_cache_friendly(
     # directly get input_sets, output_set and size_dict by using identity function as algorithm
 
 
+get_tn_info = partial(_get_path_cache_friendly, algorithm=_identity)
+
+
 # some contractor setup usages
 """
 import cotengra as ctg
@@ -513,7 +575,8 @@ opt = ctg.ReusableHyperOptimizer(
 
 def opt_reconf(inputs, output, size, **kws):
     tree = opt.search(inputs, output, size)
-    tree_r = tree.subtree_reconfigure_forest(progbar=True, num_trees=10, num_restarts=20, subtree_weight_what=("size", ))
+    tree_r = tree.subtree_reconfigure_forest(progbar=True, num_trees=10, 
+                        num_restarts=20, subtree_weight_what=("size", ))
     return tree_r.get_path()
 
 tc.set_contractor("custom", optimizer=opt_reconf)
@@ -631,6 +694,51 @@ def _base(
     return final_node
 
 
+class NodesReturn(Exception):
+    """
+    Intentionally stop execution to return a value.
+    """
+
+    def __init__(self, value_to_return: Any):
+        self.value = value_to_return
+        super().__init__(
+            f"Intentionally stopping execution to return: {value_to_return}"
+        )
+
+
+def _get_sorted_nodes(nodes: List[Any], *args: Any, **kws: Any) -> Any:
+    nodes_new = sorted(nodes, key=lambda node: getattr(node, "_stable_id_", -1))
+    raise NodesReturn(nodes_new)
+
+
+def function_nodes_capture(func: Callable[[Any], Any]) -> Callable[[Any], Any]:
+    @wraps(func)
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        with runtime_contractor(method="before"):
+            try:
+                result = func(*args, **kwargs)
+                return result
+            except NodesReturn as e:
+                return e.value
+
+    return wrapper
+
+
+@contextmanager
+def runtime_nodes_capture(key: str = "nodes") -> Iterator[Any]:
+    old_contractor = getattr(thismodule, "contractor")
+    set_contractor(method="before")
+    captured_value: Dict[str, List[tn.Node]] = {}
+    try:
+        yield captured_value
+    except NodesReturn as e:
+        captured_value[key] = e.value
+    finally:
+        for module in sys.modules:
+            if module.startswith(package_name):
+                setattr(sys.modules[module], "contractor", old_contractor)
+
+
 def custom(
     nodes: List[Any],
     optimizer: Any,
@@ -701,6 +809,16 @@ def custom_stateful(
 
 # only work for custom
 def contraction_info_decorator(algorithm: Callable[..., Any]) -> Callable[..., Any]:
+    """Decorator to add contraction information logging to an optimizer.
+
+    This decorator wraps an optimization algorithm and prints detailed information
+    about the contraction cost (FLOPs, size, write) and path finding time.
+
+    :param algorithm: The optimization algorithm to decorate.
+    :type algorithm: Callable[..., Any]
+    :return: The decorated optimization algorithm.
+    :rtype: Callable[..., Any]
+    """
     from cotengra import ContractionTree
 
     def new_algorithm(
@@ -807,6 +925,9 @@ def set_contractor(
             **kws,
         )
 
+    elif method == "before":  # a hack way to get the nodes
+        cf = _get_sorted_nodes
+
     else:
         # cf = getattr(tn.contractors, method, None)
         # if not cf: