tensorcircuit-nightly 1.2.1.dev20250725__tar.gz → 1.2.1.dev20250726__tar.gz

{tensorcircuit_nightly-1.2.1.dev20250725 → tensorcircuit_nightly-1.2.1.dev20250726}/CHANGELOG.md

@@ -10,10 +10,16 @@
 
 - Add `circuit.amplitude_before()` method to return the corresponding tensornetwork nodes.
 
+- Add `with_prob` for `stabilizercircuit.measure()`.
+
+- Add `tc.cons.function_nodes_capture` decorator and `tc.cons.runtime_nodes_capture` context manager for directly return nodes before real contraction.
+
 ### Fixed
 
 - Fix the nodes order in contraction by giving each node a global `_stable_id_`.
 
+- Fix `to_dlpack` for jax version >= 0.7.
+
 ## v1.2.1
 
 ### Fixed

{tensorcircuit_nightly-1.2.1.dev20250725 → tensorcircuit_nightly-1.2.1.dev20250726}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tensorcircuit-nightly
-Version: 1.2.1.dev20250725
+Version: 1.2.1.dev20250726
 Summary: nightly release for tensorcircuit
 Home-page: https://github.com/refraction-ray/tensorcircuit-dev
 Author: TensorCircuit Authors

{tensorcircuit_nightly-1.2.1.dev20250725 → tensorcircuit_nightly-1.2.1.dev20250726}/docs/source/advance.rst

@@ -485,6 +485,15 @@ Here's an example studying entanglement asymmetry in tilted ferromagnet states:
 Randoms, Jit, Backend Agnostic, and Their Interplay
 --------------------------------------------------------
 
+This section explains how random number generation interacts with JIT compilation and backend agnosticism in TensorCircuit. Understanding this interplay is crucial for reproducible and correct simulation results, especially when using JAX.
+
+**Key Management for Reproducibility:**
+In JAX, random number generation is deterministic and relies on explicit "keys" that manage the random state. This is different from TensorFlow or NumPy, where random states are often managed implicitly. For reproducible results and correct JIT compilation, JAX requires these keys to be passed and split explicitly.
+
+**Why Explicit Key Management?**
+When a JIT-compiled function is called multiple times with the same inputs, JAX aims to produce the same output. If random numbers were generated implicitly within a JIT-compiled function, subsequent calls would produce the same "random" numbers, which is often not the desired behavior for simulations requiring true randomness across runs. 
+Explicit key management ensures that each call to a random function, even within JIT, uses a new, distinct random state derived from a split key, thus maintaining the desired randomness and reproducibility.
+
 .. code-block:: python
 
     import tensorcircuit as tc

{tensorcircuit_nightly-1.2.1.dev20250725 → tensorcircuit_nightly-1.2.1.dev20250726}/docs/source/conf.py

@@ -20,8 +20,8 @@ sys.path.insert(0, os.path.abspath("../ext/"))
 
 # -- Project information -----------------------------------------------------
 
-project = "tensorcircuit"
-copyright = "2020, TensorCircuit Development Team. Created by Shi-Xin Zhang"
+project = "tensorcircuit-ng"
+copyright = "2020, TensorCircuit Development Team. Created by Shi-Xin Zhang."
 author = "refraction-ray"
 
 # The short X.Y version

{tensorcircuit_nightly-1.2.1.dev20250725 → tensorcircuit_nightly-1.2.1.dev20250726}/docs/source/faq.rst

@@ -4,7 +4,7 @@ Frequently Asked Questions
 What is the relation between TensorCircuit and TensorCircuit-NG?
 -------------------------------------------------------------------
 
-Both packages are created by `Shi-Xin Zhang <https://www.iop.cas.cn/rcjy/tpyjy/?id=6789>`_ (`@refraction-ray <https://github.com/refraction-ray>`_). For the history of the evolution of tensorcircuit, please refer to `history <https://github.com/tensorcircuit/tensorcircuit-ng/blob/master/HISTORY.md>`_.
+Both packages are created by `Shi-Xin Zhang <https://www.iop.cas.cn/rcjy/tpyjy/?id=6789>`_ (`@refraction-ray <https://github.com/refraction-ray>`_). For the history of the evolution of TensorCircuit-NG, please refer to `history <https://github.com/tensorcircuit/tensorcircuit-ng/blob/master/HISTORY.md>`_.
 
 From users' perspective, TensorCircuit-NG maintains full compatibility with the TensorCircuit API, enhancing it with additional features and critical bug fixes. Only TensorCircuit-NG is kept up-to-date with the fast evolving scientific computing and machine learning ecosystem in Python.
 
@@ -41,6 +41,7 @@ How can I use multiple GPUs?
 For different observables evaluation on different cards, see `example <https://github.com/tensorcircuit/tensorcircuit-ng/blob/master/examples/vqe_parallel_pmap.py>`_.
 
 For distributed simulation of one circuit on multiple cards, see `example for expectation <https://github.com/tensorcircuit/tensorcircuit-ng/blob/master/examples/slicing_auto_pmap_vqa.py>`_ and `example for MPO <https://github.com/tensorcircuit/tensorcircuit-ng/blob/master/examples/slicing_auto_pmap_mpo.py>`_.
+We also introduce a new interface for the multi-GPU tensornetwork contraction, see `example for VQE <https://github.com/tensorcircuit/tensorcircuit-ng/blob/master/examples/distributed_interface_vqe.py>`_ and `example for amplitude <https://github.com/tensorcircuit/tensorcircuit-ng/blob/master/examples/distributed_interface_amplitude.py>`_.
 
 
 When should I jit the function?

{tensorcircuit_nightly-1.2.1.dev20250725 → tensorcircuit_nightly-1.2.1.dev20250726}/docs/source/quickstart.rst

@@ -10,26 +10,25 @@ Installation
 ``pip install tensorcircuit-ng`` 
 
 is in general enough. 
-Either pip from conda or other python env managers is fine.
 
 Since there are many optional packages for various features, 
-the users may need to install more pip packages when required. 
+the users may need to install more packages by pip when required. 
 
 For Nvidia GPU,
 please refer to the GPU aware installation guide of corresponding machine learning frameworks: 
-`TensorFlow <https://www.tensorflow.org/install/gpu>`_, 
-`Jax <https://github.com/google/jax#pip-installation-gpu-cuda>`_, 
+`TensorFlow <https://www.tensorflow.org/install/pip>`_ (``pip install 'tensorflow[and-cuda]'``), 
+`Jax <https://docs.jax.dev/en/latest/installation.html#pip-installation-nvidia-gpu-cuda-installed-via-pip-easier>`_ (``pip install 'jax[cuda-12]'``), 
 or `PyTorch <https://pytorch.org/get-started/locally/>`_.
 
-Docker is also recommended (especially Linux + Nvidia GPU setup): 
+Docker is also accessible: 
 
 ``sudo docker run -it --network host --gpus all tensorcircuit/tensorcircuit``.
 
 For more details on docker setup, please refer to `docker readme <https://github.com/tensorcircuit/tensorcircuit-ng/tree/master/docker>`_.
 
 
-Overall, the installation of TensorCircuit-NG is simple, since it is purely in Python and hence very portable. 
-As long as the users can take care of the installation of ML frameworks on the corresponding system, TensorCircuit-NG will work as expected.
+Overall, the installation of TensorCircuit-NG is simple, since it is purely in Python and hence portable. 
+As long as the users can take care of the installation of ML frameworks on the corresponding operating system, TensorCircuit-NG will work as expected.
 
 To debug the installation issue or report bugs, please check the environment information by ``tc.about()``.
 
@@ -218,7 +217,7 @@ If the user has no intention to maintain the application code in a backend agnos
 Automatic Differentiation, JIT, and Vectorized Parallelism
 -------------------------------------------------------------
 
-For concepts of AD, JIT and VMAP, please refer to `Jax documentation <https://jax.readthedocs.io/en/latest/jax-101/index.html>`__ .
+For concepts of AD, JIT and VMAP, please refer to `Jax documentation <https://docs.jax.dev/en/latest/tutorials.html>`__ .
 
 The related API design in TensorCircuit-NG closely follows the functional programming design pattern in Jax with some slight differences. So we strongly recommend users learn some basics about Jax no matter which ML backend they intend to use.
 

{tensorcircuit_nightly-1.2.1.dev20250725 → tensorcircuit_nightly-1.2.1.dev20250726}/tensorcircuit/__init__.py

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

{tensorcircuit_nightly-1.2.1.dev20250725 → tensorcircuit_nightly-1.2.1.dev20250726}/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_nightly-1.2.1.dev20250725 → tensorcircuit_nightly-1.2.1.dev20250726}/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_nightly-1.2.1.dev20250725 → tensorcircuit_nightly-1.2.1.dev20250726}/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: