tensorcircuit-nightly 1.2.1.dev20250722__tar.gz → 1.2.1.dev20250724__tar.gz

{tensorcircuit_nightly-1.2.1.dev20250722 → tensorcircuit_nightly-1.2.1.dev20250724}/CHANGELOG.md

@@ -1,8 +1,19 @@
 # Change Log
 
 ## Unreleased
+
+### Added
+
 - Add `Lattice` module (`tensorcircuit.templates.lattice`) for creating and manipulating various lattice geometries, including `SquareLattice`, `HoneycombLattice`, and `CustomizeLattice`.
 
+- Add `DistributedContractor` in experimental module with new examples for fast implementation of distribution circuit simulation on jax backend.
+
+- Add `circuit.amplitude_before()` method to return the corresponding tensornetwork nodes.
+
+### Fixed
+
+- Fix the nodes order in contraction by giving each node a global `_stable_id_`.
+
 ## v1.2.1
 
 ### Fixed

{tensorcircuit_nightly-1.2.1.dev20250722 → tensorcircuit_nightly-1.2.1.dev20250724}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tensorcircuit-nightly
-Version: 1.2.1.dev20250722
+Version: 1.2.1.dev20250724
 Summary: nightly release for tensorcircuit
 Home-page: https://github.com/refraction-ray/tensorcircuit-dev
 Author: TensorCircuit Authors
@@ -60,17 +60,17 @@ Dynamic: summary
 
 <p align="center"> English | <a href="README_cn.md"> 简体中文 </a></p>
 
-TensorCircuit-NG is the next-generation open-source high-performance quantum software framework, built upon tensornetwork engines, supporting for automatic differentiation, just-in-time compiling, hardware acceleration, and vectorized parallelism, providing unified infrastructures and interfaces for quantum programming. It can compose quantum circuits, neural networks and tensor networks seamlessly with high simulation efficiency and flexibility.
+TensorCircuit-NG is the next-generation open-source high-performance quantum software framework, built upon tensornetwork engines, supporting for automatic differentiation, just-in-time compiling, hardware acceleration, vectorized parallelism and distributed training, providing unified infrastructures and interfaces for quantum programming. It can compose quantum circuits, neural networks and tensor networks seamlessly with high simulation efficiency and flexibility.
 
-TensorCircuit-NG is built on top of modern machine learning frameworks: Jax, TensorFlow, and PyTorch. It is specifically suitable for large-scale simulations of quantum-classical hybrid paradigm and variational quantum algorithms in ideal, noisy, Clifford, approximate and analog cases. It also supports quantum hardware access and provides CPU/GPU/QPU hybrid deployment solutions.
+TensorCircuit-NG is built on top of modern machine learning frameworks: Jax, TensorFlow, and PyTorch. It is specifically suitable for large-scale simulations of quantum-classical hybrid paradigm and variational quantum algorithms in ideal, noisy, Clifford, approximate, analog and fermionic cases. It also supports quantum hardware access and provides CPU/GPU/QPU hybrid deployment solutions.
 
-TensorCircuit-NG is the actively maintained official version and a [fully compatible](https://tensorcircuit-ng.readthedocs.io/en/latest/faq.html#what-is-the-relation-between-tensorcircuit-and-tensorcircuit-ng) successor to TensorCircuit with more new features (stabilizer circuit and distributed simulation) and bug fixes (support latest `numpy>2` and `qiskit>1`).
+TensorCircuit-NG is the actively maintained official version and a [fully compatible](https://tensorcircuit-ng.readthedocs.io/en/latest/faq.html#what-is-the-relation-between-tensorcircuit-and-tensorcircuit-ng) successor to TensorCircuit with more new features (stabilizer circuit, multi-card distributed simulation, etc.) and bug fixes (support latest `numpy>2` and `qiskit>1`).
 
 ## Getting Started
 
 Please begin with [Quick Start](/docs/source/quickstart.rst) in the [full documentation](https://tensorcircuit-ng.readthedocs.io/).
 
-For more information on software usage, sota algorithm implementation and engineer paradigm demonstration, please refer to 80+ [example scripts](/examples) and 30+ [tutorial notebooks](https://tensorcircuit-ng.readthedocs.io/en/latest/#tutorials). API docstrings and test cases in [tests](/tests) are also informative. One can also refer to tensorcircuit-ng [deepwiki](https://deepwiki.com/tensorcircuit/tensorcircuit-ng) generated by LLM.
+For more information on software usage, sota algorithm implementation and engineer paradigm demonstration, please refer to 80+ [example scripts](/examples) and 30+ [tutorial notebooks](https://tensorcircuit-ng.readthedocs.io/en/latest/#tutorials). API docstrings and test cases in [tests](/tests) are also informative. One can also refer to tensorcircuit-ng [deepwiki](https://deepwiki.com/tensorcircuit/tensorcircuit-ng) by LLM.
 
 For beginners, please refer to [quantum computing lectures with TC-NG](https://github.com/sxzgroup/qc_lecture) to learn both quantum computing basics and representative usage of TensorCircuit-NG.
 
@@ -170,7 +170,7 @@ The package is written in pure Python and can be obtained via pip as:
 pip install tensorcircuit-ng
 ```
 
-We recommend you install this package with tensorflow also installed as:
+We recommend you install this package with tensorflow or jax also installed as:
 
 ```python
 pip install "tensorcircuit-ng[tensorflow]"

{tensorcircuit_nightly-1.2.1.dev20250722 → tensorcircuit_nightly-1.2.1.dev20250724}/README.md

@@ -25,17 +25,17 @@
 
 <p align="center"> English | <a href="README_cn.md"> 简体中文 </a></p>
 
-TensorCircuit-NG is the next-generation open-source high-performance quantum software framework, built upon tensornetwork engines, supporting for automatic differentiation, just-in-time compiling, hardware acceleration, and vectorized parallelism, providing unified infrastructures and interfaces for quantum programming. It can compose quantum circuits, neural networks and tensor networks seamlessly with high simulation efficiency and flexibility.
+TensorCircuit-NG is the next-generation open-source high-performance quantum software framework, built upon tensornetwork engines, supporting for automatic differentiation, just-in-time compiling, hardware acceleration, vectorized parallelism and distributed training, providing unified infrastructures and interfaces for quantum programming. It can compose quantum circuits, neural networks and tensor networks seamlessly with high simulation efficiency and flexibility.
 
-TensorCircuit-NG is built on top of modern machine learning frameworks: Jax, TensorFlow, and PyTorch. It is specifically suitable for large-scale simulations of quantum-classical hybrid paradigm and variational quantum algorithms in ideal, noisy, Clifford, approximate and analog cases. It also supports quantum hardware access and provides CPU/GPU/QPU hybrid deployment solutions.
+TensorCircuit-NG is built on top of modern machine learning frameworks: Jax, TensorFlow, and PyTorch. It is specifically suitable for large-scale simulations of quantum-classical hybrid paradigm and variational quantum algorithms in ideal, noisy, Clifford, approximate, analog and fermionic cases. It also supports quantum hardware access and provides CPU/GPU/QPU hybrid deployment solutions.
 
-TensorCircuit-NG is the actively maintained official version and a [fully compatible](https://tensorcircuit-ng.readthedocs.io/en/latest/faq.html#what-is-the-relation-between-tensorcircuit-and-tensorcircuit-ng) successor to TensorCircuit with more new features (stabilizer circuit and distributed simulation) and bug fixes (support latest `numpy>2` and `qiskit>1`).
+TensorCircuit-NG is the actively maintained official version and a [fully compatible](https://tensorcircuit-ng.readthedocs.io/en/latest/faq.html#what-is-the-relation-between-tensorcircuit-and-tensorcircuit-ng) successor to TensorCircuit with more new features (stabilizer circuit, multi-card distributed simulation, etc.) and bug fixes (support latest `numpy>2` and `qiskit>1`).
 
 ## Getting Started
 
 Please begin with [Quick Start](/docs/source/quickstart.rst) in the [full documentation](https://tensorcircuit-ng.readthedocs.io/).
 
-For more information on software usage, sota algorithm implementation and engineer paradigm demonstration, please refer to 80+ [example scripts](/examples) and 30+ [tutorial notebooks](https://tensorcircuit-ng.readthedocs.io/en/latest/#tutorials). API docstrings and test cases in [tests](/tests) are also informative. One can also refer to tensorcircuit-ng [deepwiki](https://deepwiki.com/tensorcircuit/tensorcircuit-ng) generated by LLM.
+For more information on software usage, sota algorithm implementation and engineer paradigm demonstration, please refer to 80+ [example scripts](/examples) and 30+ [tutorial notebooks](https://tensorcircuit-ng.readthedocs.io/en/latest/#tutorials). API docstrings and test cases in [tests](/tests) are also informative. One can also refer to tensorcircuit-ng [deepwiki](https://deepwiki.com/tensorcircuit/tensorcircuit-ng) by LLM.
 
 For beginners, please refer to [quantum computing lectures with TC-NG](https://github.com/sxzgroup/qc_lecture) to learn both quantum computing basics and representative usage of TensorCircuit-NG.
 
@@ -135,7 +135,7 @@ The package is written in pure Python and can be obtained via pip as:
 pip install tensorcircuit-ng
 ```
 
-We recommend you install this package with tensorflow also installed as:
+We recommend you install this package with tensorflow or jax also installed as:
 
 ```python
 pip install "tensorcircuit-ng[tensorflow]"

{tensorcircuit_nightly-1.2.1.dev20250722 → tensorcircuit_nightly-1.2.1.dev20250724}/tensorcircuit/__init__.py

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

{tensorcircuit_nightly-1.2.1.dev20250722 → tensorcircuit_nightly-1.2.1.dev20250724}/tensorcircuit/basecircuit.py

@@ -441,27 +441,17 @@ class BaseCircuit(AbstractCircuit):
 
     measure = measure_jit
 
-    def amplitude(self, l: Union[str, Tensor]) -> Tensor:
+    def amplitude_before(self, l: Union[str, Tensor]) -> List[Gate]:
         r"""
-        Returns the amplitude of the circuit given the bitstring l.
+        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)`
         Note how these two are different up to a square operation.
 
-        :Example:
-
-        >>> c = tc.Circuit(2)
-        >>> c.X(0)
-        >>> c.amplitude("10")
-        array(1.+0.j, dtype=complex64)
-        >>> c.CNOT(0, 1)
-        >>> c.amplitude("11")
-        array(1.+0.j, dtype=complex64)
-
         :param l: The bitstring of 0 and 1s.
         :type l: Union[str, Tensor]
-        :return: The amplitude of the circuit.
-        :rtype: tn.Node.tensor
+        :return: The tensornetwork nodes for the amplitude of the circuit.
+        :rtype: List[Gate]
         """
         no, d_edges = self._copy()
         ms = []
@@ -502,6 +492,32 @@ class BaseCircuit(AbstractCircuit):
         no.extend(ms)
         if self.is_dm:
             no.extend(msconj)
+        return no
+
+    def amplitude(self, l: Union[str, Tensor]) -> Tensor:
+        r"""
+        Returns 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)`
+        Note how these two are different up to a square operation.
+
+        :Example:
+
+        >>> c = tc.Circuit(2)
+        >>> c.X(0)
+        >>> c.amplitude("10")
+        array(1.+0.j, dtype=complex64)
+        >>> c.CNOT(0, 1)
+        >>> c.amplitude("11")
+        array(1.+0.j, dtype=complex64)
+
+        :param l: The bitstring of 0 and 1s.
+        :type l: Union[str, Tensor]
+        :return: The amplitude of the circuit.
+        :rtype: tn.Node.tensor
+        """
+        no = self.amplitude_before(l)
+
         return contractor(no).tensor
 
     def probability(self) -> Tensor:

{tensorcircuit_nightly-1.2.1.dev20250722 → tensorcircuit_nightly-1.2.1.dev20250724}/tensorcircuit/cons.py

@@ -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"
@@ -477,39 +510,29 @@ def _identity(*args: Any, **kws: Any) -> Any:
     return args
 
 
-def _sort_tuple_list(input_list: List[Any], output_list: List[Any]) -> List[Any]:
-    sorted_elements = [(tuple(sorted(t)), i) for i, t in enumerate(input_list)]
-    sorted_elements.sort()
-    return [output_list[i] for _, i in sorted_elements]
-
-
 def _get_path_cache_friendly(
     nodes: List[tn.Node], algorithm: Any
 ) -> Tuple[List[Tuple[int, int]], List[tn.Node]]:
     nodes = list(nodes)
-    mapping_dict = {}
-    i = 0
-    for n in nodes:
-        for e in n:
-            if id(e) not in mapping_dict:
-                mapping_dict[id(e)] = get_symbol(i)
-                i += 1
-    # TODO(@refraction-ray): may be not that cache friendly, since the edge id correspondence is not that fixed?
-    input_sets = [list([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)), dtype=object))  # type: ignore
-    # nodes_new = [nodes[i] for i in order]
-    nodes_new = _sort_tuple_list(input_sets, 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 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 tn.get_subgraph_dangling(nodes_new)]
+        [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 tn.get_all_edges(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)

{tensorcircuit_nightly-1.2.1.dev20250722 → tensorcircuit_nightly-1.2.1.dev20250724}/tensorcircuit/experimental.py

@@ -3,16 +3,19 @@ Experimental features
 """
 
 from functools import partial
-from typing import Any, Callable, Optional, Tuple, Sequence, Union
+import logging
+from typing import Any, Callable, Dict, Optional, Tuple, List, Sequence, Union
 
 import numpy as np
 
-from .cons import backend, dtypestr, contractor, rdtypestr
+from .cons import backend, dtypestr, contractor, rdtypestr, get_tn_info
 from .gates import Gate
 
 Tensor = Any
 Circuit = Any
 
+logger = logging.getLogger(__name__)
+
 
 def adaptive_vmap(
     f: Callable[..., Any],
@@ -504,7 +507,7 @@ def evol_local(
     h_fun: Callable[..., Tensor],
     t: float,
     *args: Any,
-    **solver_kws: Any
+    **solver_kws: Any,
 ) -> Circuit:
     """
     ode evolution of time dependent Hamiltonian on circuit of given indices
@@ -626,3 +629,343 @@ def jax_jitted_function_load(filename: str) -> Callable[..., Any]:
 
 
 jax_func_load = jax_jitted_function_load
+
+
+PADDING_VALUE = -1
+jaxlib: Any
+ctg: Any
+
+
+class DistributedContractor:
+    """
+    A distributed tensor network contractor that parallelizes computations across multiple devices.
+
+    This class uses cotengra to find optimal contraction paths and distributes the computational
+    load across multiple devices (e.g., GPUs) for efficient tensor network calculations.
+    Particularly useful for large-scale quantum circuit simulations and variational quantum algorithms.
+
+    Example:
+        >>> def nodes_fn(params):
+        ...     c = tc.Circuit(4)
+        ...     c.rx(0, theta=params[0])
+        ...     return c.expectation_before([tc.gates.z(), [0]], reuse=False)
+        >>> dc = DistributedContractor(nodes_fn, params)
+        >>> value, grad = dc.value_and_grad(params)
+
+    :param nodes_fn: Function that takes parameters and returns a list of tensor network nodes
+    :type nodes_fn: Callable[[Tensor], List[Gate]]
+    :param params: Initial parameters used to determine the tensor network structure
+    :type params: Tensor
+    :param cotengra_options: Configuration options passed to the cotengra optimizer. Defaults to None
+    :type cotengra_options: Optional[Dict[str, Any]], optional
+    :param devices: List of devices to use. If None, uses all available local devices
+    :type devices: Optional[List[Any]], optional
+    """
+
+    def __init__(
+        self,
+        nodes_fn: Callable[[Tensor], List[Gate]],
+        params: Tensor,
+        cotengra_options: Optional[Dict[str, Any]] = None,
+        devices: Optional[List[Any]] = None,
+    ) -> None:
+        global jaxlib
+        global ctg
+
+        logger.info("Initializing DistributedContractor...")
+        import cotengra as ctg
+        import jax as jaxlib
+
+        self.nodes_fn = nodes_fn
+        if devices is None:
+            self.num_devices = jaxlib.local_device_count()
+            self.devices = jaxlib.local_devices()
+            # TODO(@refraction-ray): multi host support
+        else:
+            self.devices = devices
+            self.num_devices = len(devices)
+
+        if self.num_devices <= 1:
+            logger.info("DistributedContractor is running on a single device.")
+
+        self._params_template = params
+        self._backend = "jax"
+        self._compiled_v_fns: Dict[
+            Tuple[Callable[[Tensor], Tensor], str],
+            Callable[[Any, Tensor, Tensor], Tensor],
+        ] = {}
+        self._compiled_vg_fns: Dict[
+            Tuple[Callable[[Tensor], Tensor], str],
+            Callable[[Any, Tensor, Tensor], Tensor],
+        ] = {}
+
+        logger.info("Running cotengra pathfinder... (This may take a while)")
+        nodes = self.nodes_fn(self._params_template)
+        tn_info, _ = get_tn_info(nodes)
+        default_cotengra_options = {
+            "slicing_reconf_opts": {"target_size": 2**28},
+            "max_repeats": 128,
+            "progbar": True,
+            "minimize": "write",
+            "parallel": "auto",
+        }
+        if cotengra_options:
+            default_cotengra_options = cotengra_options
+
+        opt = ctg.ReusableHyperOptimizer(**default_cotengra_options)
+        self.tree = opt.search(*tn_info)
+        actual_num_slices = self.tree.nslices
+
+        print("\n--- Contraction Path Info ---")
+        stats = self.tree.contract_stats()
+        print(f"Path found with {actual_num_slices} slices.")
+        print(
+            f"Arithmetic Intensity (higher is better): {self.tree.arithmetic_intensity():.2f}"
+        )
+        print("flops (TFlops):", stats["flops"] / 2**40 / self.num_devices)
+        print("write (GB):", stats["write"] / 2**27 / actual_num_slices)
+        print("size (GB):", stats["size"] / 2**27)
+        print("-----------------------------\n")
+
+        slices_per_device = int(np.ceil(actual_num_slices / self.num_devices))
+        padded_size = slices_per_device * self.num_devices
+        slice_indices = np.arange(actual_num_slices)
+        padded_slice_indices = np.full(padded_size, PADDING_VALUE, dtype=np.int32)
+        padded_slice_indices[:actual_num_slices] = slice_indices
+        self.batched_slice_indices = backend.convert_to_tensor(
+            padded_slice_indices.reshape(self.num_devices, slices_per_device)
+        )
+        print(
+            f"Distributing across {self.num_devices} devices. Each device will sequentially process "
+            f"up to {slices_per_device} slices."
+        )
+
+        self._compiled_vg_fn = None
+        self._compiled_v_fn = None
+
+        logger.info("Initialization complete.")
+
+    def _get_single_slice_contraction_fn(
+        self, op: Optional[Callable[[Tensor], Tensor]] = None
+    ) -> Callable[[Any, Tensor, int], Tensor]:
+        if op is None:
+            op = backend.sum
+
+        def single_slice_contraction(
+            tree: ctg.ContractionTree, params: Tensor, slice_idx: int
+        ) -> Tensor:
+            nodes = self.nodes_fn(params)
+            _, standardized_nodes = get_tn_info(nodes)
+            input_arrays = [node.tensor for node in standardized_nodes]
+            sliced_arrays = tree.slice_arrays(input_arrays, slice_idx)
+            result = tree.contract_core(sliced_arrays, backend=self._backend)
+            return op(result)
+
+        return single_slice_contraction
+
+    def _get_device_sum_vg_fn(
+        self,
+        op: Optional[Callable[[Tensor], Tensor]] = None,
+        output_dtype: Optional[str] = None,
+    ) -> Callable[[Any, Tensor, Tensor], Tuple[Tensor, Tensor]]:
+        post_processing = lambda x: backend.real(backend.sum(x))
+        if op is None:
+            op = post_processing
+        base_fn = self._get_single_slice_contraction_fn(op=op)
+        # to ensure the output is real so that can be differentiated
+        single_slice_vg_fn = jaxlib.value_and_grad(base_fn, argnums=1)
+
+        if output_dtype is None:
+            output_dtype = rdtypestr
+
+        def device_sum_fn(
+            tree: ctg.ContractionTree, params: Tensor, slice_indices_for_device: Tensor
+        ) -> Tuple[Tensor, Tensor]:
+            def scan_body(
+                carry: Tuple[Tensor, Tensor], slice_idx: Tensor
+            ) -> Tuple[Tuple[Tensor, Tensor], None]:
+                acc_value, acc_grads = carry
+
+                def compute_and_add() -> Tuple[Tensor, Tensor]:
+                    value_slice, grads_slice = single_slice_vg_fn(
+                        tree, params, slice_idx
+                    )
+                    new_value = acc_value + value_slice
+                    new_grads = jaxlib.tree_util.tree_map(
+                        jaxlib.numpy.add, acc_grads, grads_slice
+                    )
+                    return new_value, new_grads
+
+                def do_nothing() -> Tuple[Tensor, Tensor]:
+                    return acc_value, acc_grads
+
+                return (
+                    jaxlib.lax.cond(
+                        slice_idx == PADDING_VALUE, do_nothing, compute_and_add
+                    ),
+                    None,
+                )
+
+            initial_carry = (
+                backend.cast(backend.convert_to_tensor(0.0), dtype=output_dtype),
+                jaxlib.tree_util.tree_map(lambda x: jaxlib.numpy.zeros_like(x), params),
+            )
+            (final_value, final_grads), _ = jaxlib.lax.scan(
+                scan_body, initial_carry, slice_indices_for_device
+            )
+            return final_value, final_grads
+
+        return device_sum_fn
+
+    def _get_device_sum_v_fn(
+        self,
+        op: Optional[Callable[[Tensor], Tensor]] = None,
+        output_dtype: Optional[str] = None,
+    ) -> Callable[[Any, Tensor, Tensor], Tensor]:
+        base_fn = self._get_single_slice_contraction_fn(op=op)
+        if output_dtype is None:
+            output_dtype = dtypestr
+
+        def device_sum_fn(
+            tree: ctg.ContractionTree, params: Tensor, slice_indices_for_device: Tensor
+        ) -> Tensor:
+            def scan_body(
+                carry_value: Tensor, slice_idx: Tensor
+            ) -> Tuple[Tensor, None]:
+                def compute_and_add() -> Tensor:
+                    return carry_value + base_fn(tree, params, slice_idx)
+
+                return (
+                    jaxlib.lax.cond(
+                        slice_idx == PADDING_VALUE, lambda: carry_value, compute_and_add
+                    ),
+                    None,
+                )
+
+            initial_carry = backend.cast(
+                backend.convert_to_tensor(0.0), dtype=output_dtype
+            )
+            final_value, _ = jaxlib.lax.scan(
+                scan_body, initial_carry, slice_indices_for_device
+            )
+            return final_value
+
+        return device_sum_fn
+
+    def _get_or_compile_fn(
+        self,
+        cache: Dict[
+            Tuple[Callable[[Tensor], Tensor], str],
+            Callable[[Any, Tensor, Tensor], Tensor],
+        ],
+        fn_getter: Callable[..., Any],
+        op: Optional[Callable[[Tensor], Tensor]],
+        output_dtype: Optional[str],
+    ) -> Callable[[Any, Tensor, Tensor], Tensor]:
+        """
+        Gets a compiled pmap-ed function from cache or compiles and caches it.
+
+        The cache key is a tuple of (op, output_dtype). Caution on lambda function!
+
+        Returns:
+            The compiled, pmap-ed JAX function.
+        """
+        cache_key = (op, output_dtype)
+        if cache_key not in cache:
+            device_fn = fn_getter(op=op, output_dtype=output_dtype)
+            compiled_fn = jaxlib.pmap(
+                device_fn,
+                in_axes=(
+                    None,
+                    None,
+                    0,
+                ),  # tree: broadcast, params: broadcast, indices: map
+                static_broadcasted_argnums=(0,),  # arg 0 (tree) is a static argument
+                devices=self.devices,
+            )
+            cache[cache_key] = compiled_fn  # type: ignore
+        return cache[cache_key]  # type: ignore
+
+    def value_and_grad(
+        self,
+        params: Tensor,
+        aggregate: bool = True,
+        op: Optional[Callable[[Tensor], Tensor]] = None,
+        output_dtype: Optional[str] = None,
+    ) -> Tuple[Tensor, Tensor]:
+        """
+        Calculates the value and gradient, compiling the pmap function if needed for the first call.
+
+        :param params: Parameters for the `nodes_fn` input
+        :type params: Tensor
+        :param aggregate: Whether to aggregate (sum) the results across devices, defaults to True
+        :type aggregate: bool, optional
+        :param op: Optional post-processing function for the output, defaults to None (corresponding to `backend.real`)
+            op is a cache key, so dont directly pass lambda function for op
+        :type op: Optional[Callable[[Tensor], Tensor]], optional
+        :param output_dtype: dtype str for the output of `nodes_fn`, defaults to None (corresponding to `rdtypestr`)
+        :type output_dtype: Optional[str], optional
+        """
+        compiled_vg_fn = self._get_or_compile_fn(
+            cache=self._compiled_vg_fns,
+            fn_getter=self._get_device_sum_vg_fn,
+            op=op,
+            output_dtype=output_dtype,
+        )
+
+        device_values, device_grads = compiled_vg_fn(
+            self.tree, params, self.batched_slice_indices
+        )
+
+        if aggregate:
+            total_value = backend.sum(device_values)
+            total_grad = jaxlib.tree_util.tree_map(
+                lambda x: backend.sum(x, axis=0), device_grads
+            )
+            return total_value, total_grad
+        return device_values, device_grads
+
+    def value(
+        self,
+        params: Tensor,
+        aggregate: bool = True,
+        op: Optional[Callable[[Tensor], Tensor]] = None,
+        output_dtype: Optional[str] = None,
+    ) -> Tensor:
+        """
+        Calculates the value, compiling the pmap function for the first call.
+
+        :param params: Parameters for the `nodes_fn` input
+        :type params: Tensor
+        :param aggregate: Whether to aggregate (sum) the results across devices, defaults to True
+        :type aggregate: bool, optional
+        :param op: Optional post-processing function for the output, defaults to None (corresponding to identity)
+            op is a cache key, so dont directly pass lambda function for op
+        :type op: Optional[Callable[[Tensor], Tensor]], optional
+        :param output_dtype: dtype str for the output of `nodes_fn`, defaults to None (corresponding to `dtypestr`)
+        :type output_dtype: Optional[str], optional
+        """
+        compiled_v_fn = self._get_or_compile_fn(
+            cache=self._compiled_v_fns,
+            fn_getter=self._get_device_sum_v_fn,
+            op=op,
+            output_dtype=output_dtype,
+        )
+
+        device_values = compiled_v_fn(self.tree, params, self.batched_slice_indices)
+
+        if aggregate:
+            return backend.sum(device_values)
+        return device_values
+
+    def grad(
+        self,
+        params: Tensor,
+        aggregate: bool = True,
+        op: Optional[Callable[[Tensor], Tensor]] = None,
+        output_dtype: Optional[str] = None,
+    ) -> Tensor:
+        _, grad = self.value_and_grad(
+            params, aggregate=aggregate, op=op, output_dtype=output_dtype
+        )
+        return grad

{tensorcircuit_nightly-1.2.1.dev20250722 → tensorcircuit_nightly-1.2.1.dev20250724}/tensorcircuit/quantum.py

@@ -71,7 +71,7 @@ def _reachable(nodes: List[AbstractNode]) -> List[AbstractNode]:
                 if n not in seen_nodes and n not in node_que[i + 1 :]:
                     node_que.append(n)
         i += 1
-    return seen_nodes
+    return sorted(seen_nodes, key=lambda node: getattr(node, "_stable_id_", -1))
 
 
 def reachable(
@@ -1164,7 +1164,7 @@ def tn2qop(tn_mpo: Any) -> QuOperator:
     nwires = len(tn_mpo)
     mpo = []
     for i in range(nwires):
-        mpo.append(Node(tn_mpo[i]))
+        mpo.append(Node(tn_mpo[i], name=f"mpo_{i}"))
 
     for i in range(nwires - 1):
         connect(mpo[i][1], mpo[i + 1][0])