tensorcircuit-nightly 1.2.1.dev20250725__py3-none-any.whl → 1.3.0.dev20250727__py3-none-any.whl

tensorcircuit/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "1.2.1.dev20250725"
+__version__ = "1.3.0.dev20250727"
 __author__ = "TensorCircuit Authors"
 __creator__ = "refraction-ray"
 

tensorcircuit/backends/jax_backend.py

@@ -442,7 +442,14 @@ class JaxBackend(jax_backend.JaxBackend, ExtendedBackend):  # type: ignore
     def to_dlpack(self, a: Tensor) -> Any:
         import jax.dlpack
 
-        return jax.dlpack.to_dlpack(a)
+        try:
+            return jax.dlpack.to_dlpack(a)  # type: ignore
+        except AttributeError:  # jax >v0.7
+            # jax.dlpack.to_dlpack was deprecated in JAX v0.6.0 and removed in JAX v0.7.0.
+            # Please use the newer DLPack API based on __dlpack__ and __dlpack_device__ instead.
+            # Typically, you can pass a JAX array directly to the `from_dlpack` function of
+            # another framework without using `to_dlpack`.
+            return a.__dlpack__()
 
     def set_random_state(
         self, seed: Optional[Union[int, PRNGKeyArray]] = None, get_only: bool = False

tensorcircuit/circuit.py

@@ -231,6 +231,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 +342,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

tensorcircuit/cons.py

@@ -516,8 +516,8 @@ def _get_path_cache_friendly(
     nodes = list(nodes)
 
     nodes_new = sorted(nodes, key=lambda node: getattr(node, "_stable_id_", -1))
-    if isinstance(algorithm, list):
-        return algorithm, nodes_new
+    # if isinstance(algorithm, list):
+    #     return algorithm, [nodes_new]
 
     all_edges = tn.get_all_edges(nodes_new)
     all_edges_sorted = sorted_edges(all_edges)
@@ -693,6 +693,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,
@@ -763,6 +808,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(
@@ -869,6 +924,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: