tensorcircuit-nightly 1.4.0.dev20251016__py3-none-any.whl → 1.4.0.dev20251128__py3-none-any.whl

tensorcircuit/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "1.4.0.dev20251016"
+__version__ = "1.4.0.dev20251128"
 __author__ = "TensorCircuit Authors"
 __creator__ = "refraction-ray"
 

tensorcircuit/experimental.py

@@ -7,6 +7,10 @@ Experimental features
 from functools import partial
 import logging
 from typing import Any, Callable, Dict, Optional, Tuple, List, Sequence, Union
+import pickle
+import uuid
+import time
+import os
 
 import numpy as np
 
@@ -489,6 +493,229 @@ jax_func_load = jax_jitted_function_load
 PADDING_VALUE = -1
 jaxlib: Any
 ctg: Any
+Mesh: Any
+NamedSharding: Any
+P: Any
+
+
+def broadcast_py_object(obj: Any, shared_dir: Optional[str] = None) -> Any:
+    """
+    Broadcast a picklable Python object from process 0 to all other processes,
+    with fallback mechanism from gRPC to file system based approach.
+
+    This function first attempts to use gRPC-based broadcast. If that fails due to
+    pickling issues, it falls back to a file system based approach that is more robust.
+
+    :param obj: The Python object to broadcast. It must be picklable.
+                This object should exist on process 0 and can be None on others.
+    :type obj: Any
+    :param shared_dir: Directory path for shared file system broadcast fallback.
+                       If None, uses current directory. Only used in fallback mode.
+    :type shared_dir: Optional[str], optional
+    :return: The broadcasted object, now present on all processes.
+    :rtype: Any
+    """
+    import jax
+    from jax.experimental import multihost_utils
+
+    try:
+        result = broadcast_py_object_jax(obj)
+        return result
+
+    except pickle.UnpicklingError as e:
+        # This block is executed if any process fails during the gRPC attempt.
+
+        multihost_utils.sync_global_devices("grpc_broadcast_failed_fallback_sync")
+
+        if jax.process_index() == 0:
+            border = "=" * 80
+            logger.warning(
+                "\n%s\nJAX gRPC broadcast failed with error: %s\n"
+                "--> Falling back to robust Shared File System broadcast method.\n%s",
+                border,
+                e,
+                border,
+            )
+
+        return broadcast_py_object_fs(obj, shared_dir)
+
+
+def broadcast_py_object_jax(obj: Any) -> Any:
+    """
+    Broadcast a picklable Python object from process 0 to all other processes
+    within jax ditribution system.
+
+    This function uses a two-step broadcast: first the size, then the data.
+    This is necessary because `broadcast_one_to_all` requires the same
+    shaped array on all hosts.
+
+    :param obj: The Python object to broadcast. It must be picklable.
+             This object should exist on process 0 and can be None on others.
+
+    :return: The broadcasted object, now present on all processes.
+    """
+    import jax as jaxlib
+    import pickle
+    from jax.experimental import multihost_utils
+
+    # Serialize to bytes on process 0, empty bytes on others
+    if jaxlib.process_index() == 0:
+        if obj is None:
+            raise ValueError("Object to broadcast from process 0 cannot be None.")
+        data = pickle.dumps(obj)
+        logger.info(
+            f"--- Size of object to be broadcast: {len(data) / 1024**2:.3f} MB ---"
+        )
+
+    else:
+        data = b""
+
+    # Step 1: Broadcast the length of the serialized data.
+    # We send a single-element int32 array.
+    length = np.array([len(data)], dtype=np.int32)
+    length = multihost_utils.broadcast_one_to_all(length)
+
+    length = int(length[0])  # type: ignore
+
+    # Step 2: Broadcast the actual data.
+    # Convert byte string to a uint8 array for broadcasting.
+    send_arr_uint8 = np.frombuffer(data, dtype=np.uint8)
+    padded_length = (length + 3) // 4 * 4
+    if send_arr_uint8.size < padded_length:
+        send_arr_uint8 = np.pad(  #  type: ignore
+            send_arr_uint8, (0, padded_length - send_arr_uint8.size), mode="constant"
+        )
+    send_arr_int32 = send_arr_uint8.astype(np.int32)
+    # send_arr_int32 = jaxlib.numpy.array(send_arr_int32, dtype=np.int32)
+    send_arr_int32 = jaxlib.device_put(send_arr_int32)
+
+    jaxlib.experimental.multihost_utils.sync_global_devices("bulk_before")
+
+    received_arr = multihost_utils.broadcast_one_to_all(send_arr_int32)
+
+    received_arr = np.array(received_arr)
+    received_arr_uint8 = received_arr.astype(np.uint8)
+
+    # Step 3: Reconstruct the object from the received bytes.
+    # Convert the NumPy array back to bytes, truncate any padding, and unpickle.
+    received_data = received_arr_uint8[:length].tobytes()
+    # if jaxlib.process_index() == 0:
+    #     logger.info(f"Broadcasted object {obj}")
+    return pickle.loads(received_data)
+
+
+def broadcast_py_object_fs(
+    obj: Any, shared_dir: Optional[str] = None, timeout_seconds: int = 300
+) -> Any:
+    """
+    Broadcast a picklable Python object from process 0 to all other processes
+    using a shared file system approach.
+
+    This is a fallback method when gRPC-based broadcast fails. It uses UUID-based
+    file communication to share objects between processes through a shared file system.
+
+    :param obj: The Python object to broadcast. Must be picklable.
+                Should exist on process 0, can be None on others.
+    :type obj: Any
+    :param shared_dir: Directory path for shared file system communication.
+                       If None, uses current directory.
+    :type shared_dir: Optional[str], optional
+    :param timeout_seconds: Maximum time to wait for file operations before timing out.
+                            Defaults to 300 seconds.
+    :type timeout_seconds: int, optional
+    :return: The broadcasted object, now present on all processes.
+    :rtype: Any
+    """
+    # to_avoid very subtle bugs for broadcast tree_data on A800 clusters
+    import jax
+    from jax.experimental import multihost_utils
+
+    if shared_dir is None:
+        shared_dir = "."
+    if jax.process_index() == 0:
+        os.makedirs(shared_dir, exist_ok=True)
+
+    id_comm_path = os.path.join(shared_dir, f".broadcast_temp_12318")
+    transfer_id = ""
+
+    if jax.process_index() == 0:
+        transfer_id = str(uuid.uuid4())
+        # print(f"[Process 0] Generated unique transfer ID: {transfer_id}", flush=True)
+        with open(id_comm_path, "w") as f:
+            f.write(transfer_id)
+
+    multihost_utils.sync_global_devices("fs_broadcast_id_written")
+
+    if jax.process_index() != 0:
+        start_time = time.time()
+        while not os.path.exists(id_comm_path):
+            time.sleep(0.1)
+            if time.time() - start_time > timeout_seconds:
+                raise TimeoutError(
+                    f"Process {jax.process_index()} timed out waiting for ID file: {id_comm_path}"
+                )
+        with open(id_comm_path, "r") as f:
+            transfer_id = f.read()
+
+    multihost_utils.sync_global_devices("fs_broadcast_id_read")
+    if jax.process_index() == 0:
+        try:
+            os.remove(id_comm_path)
+        except OSError:
+            pass  # 如果文件已被其他进程快速清理，忽略错误
+
+    # 定义本次传输使用的数据文件和标志文件路径
+    data_path = os.path.join(shared_dir, f"{transfer_id}.data")
+    done_path = os.path.join(shared_dir, f"{transfer_id}.done")
+
+    result_obj = None
+
+    if jax.process_index() == 0:
+        if obj is None:
+            raise ValueError("None cannot be broadcasted.")
+
+        # print(f"[Process 0] Pickling object...", flush=True)
+        pickled_data = pickle.dumps(obj)
+        logger.info(
+            f"[Process 0] Writing {len(pickled_data) / 1024**2:.3f} MB to {data_path}"
+        )
+        with open(data_path, "wb") as f:
+            f.write(pickled_data)
+
+        with open(done_path, "w") as f:
+            pass
+        logger.info(f"[Process 0] Write complete.")
+        result_obj = obj
+    else:
+        # print(f"[Process {jax.process_index()}] Waiting for done file: {done_path}", flush=True)
+        start_time = time.time()
+        while not os.path.exists(done_path):
+            time.sleep(0.1)
+            if time.time() - start_time > timeout_seconds:
+                raise TimeoutError(
+                    f"Process {jax.process_index()} timed out waiting for done file: {done_path}"
+                )
+
+        # print(f"[Process {jax.process_index()}] Done file found. Reading data from {data_path}", flush=True)
+        with open(data_path, "rb") as f:
+            pickled_data = f.read()
+
+        result_obj = pickle.loads(pickled_data)
+        logger.info(f"[Process {jax.process_index()}] Object successfully loaded.")
+
+    multihost_utils.sync_global_devices("fs_broadcast_read_complete")
+
+    if jax.process_index() == 0:
+        try:
+            os.remove(data_path)
+            os.remove(done_path)
+            # print(f"[Process 0] Cleaned up temporary files for transfer {transfer_id}.", flush=True)
+        except OSError as e:
+            logger.info(
+                f"[Process 0]: Failed to clean up temporary files: {e}",
+            )
+
+    return result_obj
 
 
 class DistributedContractor:
@@ -513,8 +740,10 @@ class DistributedContractor:
     :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
+    :param devices: List of devices to use. If None, uses all available devices
     :type devices: Optional[List[Any]], optional
+    :param mesh: Mesh object to use for distributed computation. If None, uses all available devices
+    :type mesh: Optional[Any], optional
     """
 
     def __init__(
@@ -522,28 +751,39 @@ class DistributedContractor:
         nodes_fn: Callable[[Tensor], List[Gate]],
         params: Tensor,
         cotengra_options: Optional[Dict[str, Any]] = None,
-        devices: Optional[List[Any]] = None,
+        devices: Optional[List[Any]] = None,  # backward compatibility
+        mesh: Optional[Any] = None,
+        tree_data: Optional[Dict[str, Any]] = None,
     ) -> None:
         global jaxlib
         global ctg
+        global Mesh
+        global NamedSharding
+        global P
 
         logger.info("Initializing DistributedContractor...")
         import cotengra as ctg
+        from cotengra import ContractionTree
         import jax as jaxlib
+        from jax.sharding import Mesh, NamedSharding, PartitionSpec as P
 
         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
+        if mesh is not None:
+            self.mesh = mesh
+        elif devices is not None:
+            self.mesh = Mesh(devices, axis_names=("devices",))
         else:
-            self.devices = devices
-            self.num_devices = len(devices)
+            self.mesh = Mesh(jaxlib.devices(), axis_names=("devices",))
+        self.num_devices = len(self.mesh.devices)
 
         if self.num_devices <= 1:
             logger.info("DistributedContractor is running on a single device.")
 
         self._params_template = params
+        self.params_sharding = jaxlib.tree_util.tree_map(
+            lambda x: NamedSharding(self.mesh, P(*((None,) * x.ndim))),
+            self._params_template,
+        )
         self._backend = "jax"
         self._compiled_v_fns: Dict[
             Tuple[Callable[[Tensor], Tensor], str],
@@ -555,41 +795,68 @@ class DistributedContractor:
         ] = {}
 
         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
+        if tree_data is None:
+            if params is None:
+                raise ValueError("Please provide specific circuit parameters array.")
+            if jaxlib.process_index() == 0:
+                logger.info("Process 0: Running cotengra pathfinder...")
+                tree_data = self._get_tree_data(
+                    self.nodes_fn, self._params_template, cotengra_options  # type: ignore
+                )
 
-        opt = ctg.ReusableHyperOptimizer(**default_cotengra_options)
-        self.tree = opt.search(*tn_info)
-        actual_num_slices = self.tree.nslices
+            # Step 2: Use the robust helper function to broadcast the tree object.
+            # Process 0 sends its computed `tree_object`.
+            # Other processes send `None`, but receive the object from process 0.
 
-        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}"
+            if jaxlib.process_count() > 1:
+                # self.tree = broadcast_py_object(tree_object)
+                jaxlib.experimental.multihost_utils.sync_global_devices("tree_before")
+                logger.info(
+                    f"Process {jaxlib.process_index()}: Synchronizing contraction path..."
+                )
+                tree_data = broadcast_py_object(tree_data)
+                jaxlib.experimental.multihost_utils.sync_global_devices("tree_after")
+        else:
+            logger.info("Using pre-computed contraction path.")
+        if tree_data is None:
+            raise ValueError("Contraction path data is missing.")
+
+        self.tree = ContractionTree.from_path(
+            inputs=tree_data["inputs"],
+            output=tree_data["output"],
+            size_dict=tree_data["size_dict"],
+            path=tree_data["path"],
         )
-        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")
+
+        # Restore slicing information
+        for ind, _ in tree_data["sliced_inds"].items():
+            self.tree.remove_ind_(ind)
+
+        logger.info(
+            f"Process {jaxlib.process_index()}: Contraction path successfully synchronized."
+        )
+        actual_num_slices = self.tree.nslices
+
+        self._report_tree_info()
 
         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)
+
+        # Reshape for distribution and define the sharding rule
+        batched_indices = padded_slice_indices.reshape(
+            self.num_devices, slices_per_device
         )
+        # Sharding rule: split the first axis (the one for devices) across the 'devices' mesh axis
+        self.sharding = NamedSharding(self.mesh, P("devices", None))
+        # Place the tensor on devices according to the rule
+        self.batched_slice_indices = jaxlib.device_put(batched_indices, self.sharding)
+
+        # 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."
@@ -600,6 +867,90 @@ class DistributedContractor:
 
         logger.info("Initialization complete.")
 
+    def _report_tree_info(self) -> None:
+        print("\n--- Contraction Path Info ---")
+        actual_num_slices = self.tree.nslices
+        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")
+
+    @staticmethod
+    def _get_tree_data(
+        nodes_fn: Callable[[Tensor], List[Gate]],
+        params: Tensor,
+        cotengra_options: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, Any]:
+        global ctg
+
+        import cotengra as ctg
+
+        local_cotengra_options = (cotengra_options or {}).copy()
+
+        nodes = nodes_fn(params)
+        tn_info, _ = get_tn_info(nodes)
+        default_cotengra_options = {
+            "slicing_reconf_opts": {"target_size": 2**28},
+            "max_repeats": 128,
+            "minimize": "write",
+            "parallel": "auto",
+            "progbar": True,
+        }
+        default_cotengra_options.update(local_cotengra_options)
+
+        opt = ctg.ReusableHyperOptimizer(**default_cotengra_options)
+        tree_object = opt.search(*tn_info)
+        tree_data = {
+            "inputs": tree_object.inputs,
+            "output": tree_object.output,
+            "size_dict": tree_object.size_dict,
+            "path": tree_object.get_path(),
+            "sliced_inds": tree_object.sliced_inds,
+        }
+        return tree_data
+
+    @staticmethod
+    def find_path(
+        nodes_fn: Callable[[Tensor], Tensor],
+        params: Tensor,
+        cotengra_options: Optional[Dict[str, Any]] = None,
+        filepath: Optional[str] = None,
+    ) -> None:
+        tree_data = DistributedContractor._get_tree_data(
+            nodes_fn, params, cotengra_options
+        )
+        if filepath is not None:
+            with open(filepath, "wb") as f:
+                pickle.dump(tree_data, f)
+            logger.info(f"Contraction path data successfully saved to '{filepath}'.")
+
+    @classmethod
+    def from_path(
+        cls,
+        filepath: str,
+        nodes_fn: Callable[[Tensor], List[Gate]],
+        devices: Optional[List[Any]] = None,  # backward compatibility
+        mesh: Optional[Any] = None,
+        params: Any = None,
+    ) -> "DistributedContractor":
+        with open(filepath, "rb") as f:
+            tree_data = pickle.load(f)
+
+        # Each process loads the file independently. No broadcast is needed.
+        # We pass the loaded `tree_data` directly to __init__ to trigger the second workflow.
+        return cls(
+            nodes_fn=nodes_fn,
+            params=params,
+            mesh=mesh,
+            devices=devices,
+            tree_data=tree_data,
+        )
+
     def _get_single_slice_contraction_fn(
         self, op: Optional[Callable[[Tensor], Tensor]] = None
     ) -> Callable[[Any, Tensor, int], Tensor]:
@@ -716,6 +1067,7 @@ class DistributedContractor:
         fn_getter: Callable[..., Any],
         op: Optional[Callable[[Tensor], Tensor]],
         output_dtype: Optional[str],
+        is_grad_fn: bool,
     ) -> Callable[[Any, Tensor, Tensor], Tensor]:
         """
         Gets a compiled pmap-ed function from cache or compiles and caches it.
@@ -728,15 +1080,60 @@ class DistributedContractor:
         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,
+
+            def global_aggregated_fn(
+                tree: Any, params: Any, batched_slice_indices: Tensor
+            ) -> Any:
+                # Use jax.vmap to apply the per-device function across the sharded data.
+                # vmap maps `device_fn` over the first axis (0) of `batched_slice_indices`.
+                # `tree` and `params` are broadcasted (in_axes=None) to each call.
+                vmapped_device_fn = jaxlib.vmap(
+                    device_fn, in_axes=(None, None, 0), out_axes=0
+                )
+                device_results = vmapped_device_fn(tree, params, batched_slice_indices)
+
+                # Now, `device_results` is a sharded PyTree (one result per device).
+                # We aggregate them using jnp.sum, which JAX automatically compiles
+                # into a cross-device AllReduce operation.
+
+                if is_grad_fn:
+                    # `device_results` is a (value, grad) tuple of sharded arrays
+                    device_values, device_grads = device_results
+
+                    # Replace psum with jnp.sum
+                    global_value = jaxlib.numpy.sum(device_values, axis=0)
+                    global_grad = jaxlib.tree_util.tree_map(
+                        lambda g: jaxlib.numpy.sum(g, axis=0), device_grads
+                    )
+                    return global_value, global_grad
+                else:
+                    # `device_results` is just the sharded values
+                    return jaxlib.numpy.sum(device_results, axis=0)
+
+            #  Compile the global function with jax.jit and specify shardings.
+            # `params` are replicated (available everywhere).
+
+            in_shardings = (self.params_sharding, self.sharding)
+
+            if is_grad_fn:
+                # Returns (value, grad), so out_sharding must be a 2-tuple.
+                # `value` is a replicated scalar -> P()
+                sharding_for_value = NamedSharding(self.mesh, P())
+                # `grad` is a replicated PyTree with the same structure as params.
+                sharding_for_grad = self.params_sharding
+                out_shardings = (sharding_for_value, sharding_for_grad)
+            else:
+                # Returns a single scalar value -> P()
+                out_shardings = NamedSharding(self.mesh, P())
+
+            compiled_fn = jaxlib.jit(
+                global_aggregated_fn,
+                # `tree` is a static argument, its value is compiled into the function.
+                static_argnums=(0,),
+                # Specify how inputs are sharded.
+                in_shardings=in_shardings,
+                # Specify how the output should be sharded.
+                out_shardings=out_shardings,
             )
             cache[cache_key] = compiled_fn  # type: ignore
         return cache[cache_key]  # type: ignore
@@ -744,7 +1141,7 @@ class DistributedContractor:
     def value_and_grad(
         self,
         params: Tensor,
-        aggregate: bool = True,
+        # aggregate: bool = True,
         op: Optional[Callable[[Tensor], Tensor]] = None,
         output_dtype: Optional[str] = None,
     ) -> Tuple[Tensor, Tensor]:
@@ -753,8 +1150,6 @@ class DistributedContractor:
 
         :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
@@ -766,24 +1161,18 @@ class DistributedContractor:
             fn_getter=self._get_device_sum_vg_fn,
             op=op,
             output_dtype=output_dtype,
+            is_grad_fn=True,
         )
 
-        device_values, device_grads = compiled_vg_fn(
+        total_value, total_grad = 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
+        return total_value, total_grad
 
     def value(
         self,
         params: Tensor,
-        aggregate: bool = True,
+        # aggregate: bool = True,
         op: Optional[Callable[[Tensor], Tensor]] = None,
         output_dtype: Optional[str] = None,
     ) -> Tensor:
@@ -792,8 +1181,6 @@ class DistributedContractor:
 
         :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
@@ -805,22 +1192,17 @@ class DistributedContractor:
             fn_getter=self._get_device_sum_v_fn,
             op=op,
             output_dtype=output_dtype,
+            is_grad_fn=False,
         )
 
-        device_values = compiled_v_fn(self.tree, params, self.batched_slice_indices)
-
-        if aggregate:
-            return backend.sum(device_values)
-        return device_values
+        total_value = compiled_v_fn(self.tree, params, self.batched_slice_indices)
+        return total_value
 
     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
-        )
+        _, grad = self.value_and_grad(params, op=op, output_dtype=output_dtype)
         return grad

tensorcircuit/interfaces/tensortrans.py

@@ -132,13 +132,17 @@ def general_args_to_backend(
         target_backend = backend
     elif isinstance(target_backend, str):
         target_backend = get_backend(target_backend)
+    try:
+        t = backend.tree_map(target_backend.from_dlpack, caps)
+    except TypeError:
+        t = backend.tree_map(target_backend.from_dlpack, args)
+
     if dtype is None:
-        return backend.tree_map(target_backend.from_dlpack, caps)
+        return t
     if isinstance(dtype, str):
         leaves, treedef = backend.tree_flatten(args)
         dtype = [dtype for _ in range(len(leaves))]
         dtype = backend.tree_unflatten(treedef, dtype)
-    t = backend.tree_map(target_backend.from_dlpack, caps)
     t = backend.tree_map(target_backend.cast, t, dtype)
     return t
 

tensorcircuit/interfaces/torch.py

@@ -69,12 +69,14 @@ def torch_interface(
         @staticmethod
         def forward(ctx: Any, *x: Any) -> Any:  # type: ignore
             # ctx.xdtype = [xi.dtype for xi in x]
-            ctx.xdtype = backend.tree_map(lambda s: s.dtype, x)
+            ctx.save_for_backward(*x)
+            x_detached = backend.tree_map(lambda s: s.detach(), x)
+            ctx.xdtype = backend.tree_map(lambda s: s.dtype, x_detached)
             # (x, )
             if len(ctx.xdtype) == 1:
                 ctx.xdtype = ctx.xdtype[0]
-            ctx.device = (backend.tree_flatten(x)[0][0]).device
-            x = general_args_to_backend(x, enable_dlpack=enable_dlpack)
+            ctx.device = (backend.tree_flatten(x_detached)[0][0]).device
+            x = general_args_to_backend(x_detached, enable_dlpack=enable_dlpack)
             y = fun(*x)
             ctx.ydtype = backend.tree_map(lambda s: s.dtype, y)
             if len(x) == 1:
@@ -88,6 +90,9 @@ def torch_interface(
 
         @staticmethod
         def backward(ctx: Any, *grad_y: Any) -> Any:
+            x = ctx.saved_tensors
+            x_detached = backend.tree_map(lambda s: s.detach(), x)
+            x_backend = general_args_to_backend(x_detached, enable_dlpack=enable_dlpack)
             if len(grad_y) == 1:
                 grad_y = grad_y[0]
             grad_y = backend.tree_map(lambda s: s.contiguous(), grad_y)
@@ -96,7 +101,12 @@ def torch_interface(
             )
             # grad_y = general_args_to_numpy(grad_y)
             # grad_y = numpy_args_to_backend(grad_y, dtype=ctx.ydtype)  # backend.dtype
-            _, g = vjp_fun(ctx.x, grad_y)
+            if len(x_backend) == 1:
+                x_backend_for_vjp = x_backend[0]
+            else:
+                x_backend_for_vjp = x_backend
+
+            _, g = vjp_fun(x_backend_for_vjp, grad_y)
             # a redundency due to current vjp API
 
             r = general_args_to_backend(

{tensorcircuit_nightly-1.4.0.dev20251016.dist-info → tensorcircuit_nightly-1.4.0.dev20251128.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tensorcircuit-nightly
-Version: 1.4.0.dev20251016
+Version: 1.4.0.dev20251128
 Summary: High performance unified quantum computing framework for the NISQ era
 Author-email: TensorCircuit Authors <znfesnpbh@gmail.com>
 License-Expression: Apache-2.0
@@ -70,7 +70,7 @@ TensorCircuit-NG is the only actively maintained official version and a [fully c
 
 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 90+ [example scripts](/examples) and 40+ [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 AI-native docs for tensorcircuit-ng: [Devin Deepwiki](https://deepwiki.com/tensorcircuit/tensorcircuit-ng) and [Context7 MCP](https://context7.com/tensorcircuit/tensorcircuit-ng).
+For more information on software usage, sota algorithm implementation and engineer paradigm demonstration, please refer to 100+ [example scripts](/examples) and 40+ [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 AI-native docs for tensorcircuit-ng: [Devin Deepwiki](https://deepwiki.com/tensorcircuit/tensorcircuit-ng) and [Context7 MCP](https://context7.com/tensorcircuit/tensorcircuit-ng).
 
 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.
 
@@ -347,6 +347,7 @@ TensorCircuit-NG is open source, released under the Apache License, Version 2.0.
       <td align="center" valign="top" width="16.66%"><a href="https://github.com/Charlespkuer"><img src="https://avatars.githubusercontent.com/u/112697147?v=4?s=100" width="100px;" alt="Huang"/><br /><sub><b>Huang</b></sub></a><br /><a href="https://github.com/tensorcircuit/tensorcircuit-ng/commits?author=Charlespkuer" title="Code">💻</a> <a href="#example-Charlespkuer" title="Examples">💡</a> <a href="https://github.com/tensorcircuit/tensorcircuit-ng/commits?author=Charlespkuer" title="Tests">⚠️</a></td>
       <td align="center" valign="top" width="16.66%"><a href="https://github.com/Huang-Xu-Yang"><img src="https://avatars.githubusercontent.com/u/227286661?v=4?s=100" width="100px;" alt="Huang-Xu-Yang"/><br /><sub><b>Huang-Xu-Yang</b></sub></a><br /><a href="https://github.com/tensorcircuit/tensorcircuit-ng/commits?author=Huang-Xu-Yang" title="Code">💻</a> <a href="https://github.com/tensorcircuit/tensorcircuit-ng/commits?author=Huang-Xu-Yang" title="Tests">⚠️</a></td>
       <td align="center" valign="top" width="16.66%"><a href="https://github.com/WeiguoMa"><img src="https://avatars.githubusercontent.com/u/108172530?v=4?s=100" width="100px;" alt="Weiguo_M"/><br /><sub><b>Weiguo_M</b></sub></a><br /><a href="https://github.com/tensorcircuit/tensorcircuit-ng/commits?author=WeiguoMa" title="Code">💻</a> <a href="https://github.com/tensorcircuit/tensorcircuit-ng/commits?author=WeiguoMa" title="Tests">⚠️</a> <a href="#example-WeiguoMa" title="Examples">💡</a> <a href="#tutorial-WeiguoMa" title="Tutorials">✅</a></td>
+      <td align="center" valign="top" width="16.66%"><a href="https://github.com/QuiXamii"><img src="https://avatars.githubusercontent.com/u/136054857?v=4?s=100" width="100px;" alt="Qixiang WANG"/><br /><sub><b>Qixiang WANG</b></sub></a><br /><a href="#example-QuiXamii" title="Examples">💡</a></td>
     </tr>
   </tbody>
 </table>

{tensorcircuit_nightly-1.4.0.dev20251016.dist-info → tensorcircuit_nightly-1.4.0.dev20251128.dist-info}/RECORD

@@ -1,4 +1,4 @@
-tensorcircuit/__init__.py,sha256=A-DZjfxQbfOIeoGVNLbZOdx2tCmjBL--rIFkYud-3lk,2160
+tensorcircuit/__init__.py,sha256=FDfcQl270QZ832T4jvsAqjx4s_1pGhl1H8m6FOx212g,2160
 tensorcircuit/about.py,sha256=DazTswU2nAwOmASTaDII3L04PVtaQ7oiWPty5YMI3Wk,5267
 tensorcircuit/abstractcircuit.py,sha256=uDRgaDeH_Ym-6_ZEOZwvxHIDycVLHkGZv4zfaIgaEnc,44235
 tensorcircuit/analogcircuit.py,sha256=4BzIC631MZ2m05CXuk2T6HQ8RTmHBE6NszaOLuxmlEc,15639
@@ -8,7 +8,7 @@ tensorcircuit/channels.py,sha256=CFQxWI-JmkIxexslCBdjp_RSxUbHs6eAJv4LvlXXXCY,286
 tensorcircuit/circuit.py,sha256=lETz1SvUh_60ZMFtvSPMWOF6zWMMyQU4TyB_VwhkVHM,40027
 tensorcircuit/cons.py,sha256=V0wjevtDkESCIWMJaysgPVorQlPAIT0vtRWvIZkEWcE,33065
 tensorcircuit/densitymatrix.py,sha256=C8Q2fHXZ78S9ZaPqCIKl6_v_sILqbBgqBOUYUQ1QmFI,15020
-tensorcircuit/experimental.py,sha256=TGK4FaS6TS_ZhtjcIZgYVuAkGdRW50LN0DdXp-h4bos,29906
+tensorcircuit/experimental.py,sha256=ao7O40vaEzgS7Mbrg2AOZ3rPTdrSW7EbAWKjwqF6MPo,44975
 tensorcircuit/fgs.py,sha256=J1TjAiiqZk9KO1xYX_V0xsgKlYZaUQ7Enm4s5zkRM50,49514
 tensorcircuit/gates.py,sha256=9x1VTEpZWz-FoWVM_YznoU1dbFzXnfXIEJQQVec-2Ko,30504
 tensorcircuit/keras.py,sha256=nMSuu9uZy7haWwuen1g_6GFVwYIirtX9IvejDyoH33M,10129
@@ -71,8 +71,8 @@ tensorcircuit/interfaces/jax.py,sha256=q_nay20gcrPRyY2itvcOtkCjqtvcC4qotbvrgm2a3
 tensorcircuit/interfaces/numpy.py,sha256=T7h64dG9e5xDG0KVOy9O8TXyrt5RWRnTWN9iXf3aGyY,1439
 tensorcircuit/interfaces/scipy.py,sha256=_P2IeqvJiO7cdjTzNCIAFm8Y56Wd3j3jGmWUeeQ1Fw8,3402
 tensorcircuit/interfaces/tensorflow.py,sha256=U4hZjm-yWxOJ5tqmffk8-tNvOkAltYBJ8Z6jYwOtTaM,3355
-tensorcircuit/interfaces/tensortrans.py,sha256=oUxIVpXfANZVRXfPjiGJDzFPiszfBsiY40ydh0BaELE,10364
-tensorcircuit/interfaces/torch.py,sha256=13IFGmWUFoWiSzKAzwp2EkOSxgiwN_oUFxjQb36gimo,5149
+tensorcircuit/interfaces/tensortrans.py,sha256=YHTKIINjXE085fqO_AfUJGE-t3OThUH8csk5PFK7Dig,10414
+tensorcircuit/interfaces/torch.py,sha256=J04-bguSvJOiV-uhNVzMH28Pl-pDYPnTOv7wNm-QRZM,5633
 tensorcircuit/results/__init__.py,sha256=3kkIvmjLYQd5ff-emY8l82rpv9mwMZdM2kTLZ9sNfA4,89
 tensorcircuit/results/counts.py,sha256=gJ9x2D09wSZ8bwLB5ZR9lyx-bg6AAoz6JDr9cDAb83w,7267
 tensorcircuit/results/readout_mitigation.py,sha256=dVpNvtFZe7n_fDVczKcqYPEepu3fV2qK3u-SfOpTf68,31746
@@ -89,8 +89,8 @@ tensorcircuit/templates/graphs.py,sha256=cPYrxjoem0xZ-Is9dZKAvEzWZL_FejfIRiCEOTA
 tensorcircuit/templates/hamiltonians.py,sha256=Guvqqi-V47w8xeZDmca4_mU4mW9V4c3AplsBOrRtxFo,6308
 tensorcircuit/templates/lattice.py,sha256=IvFyNgsFMfj82g-tpJraI3lMbI-EIZ0Cghq9v7tZ6Wg,72851
 tensorcircuit/templates/measurements.py,sha256=pzc5Aa9S416Ilg4aOY77Z6ZhUlYcXnAkQNQFTuHjFFs,10943
-tensorcircuit_nightly-1.4.0.dev20251016.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-tensorcircuit_nightly-1.4.0.dev20251016.dist-info/METADATA,sha256=REXiOle-3QQShQlMCkkEI6MBM_7MhAFTCOkbPR-sIh0,38283
-tensorcircuit_nightly-1.4.0.dev20251016.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-tensorcircuit_nightly-1.4.0.dev20251016.dist-info/top_level.txt,sha256=9dcuK5488dWpVauYz8cdvx743z_La1h7zIQCsEEgu7o,14
-tensorcircuit_nightly-1.4.0.dev20251016.dist-info/RECORD,,
+tensorcircuit_nightly-1.4.0.dev20251128.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+tensorcircuit_nightly-1.4.0.dev20251128.dist-info/METADATA,sha256=XvzlNlaR6TyBER37rkHdFO_FCUIAbxTKlr_ZcvQbVZ0,38586
+tensorcircuit_nightly-1.4.0.dev20251128.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+tensorcircuit_nightly-1.4.0.dev20251128.dist-info/top_level.txt,sha256=9dcuK5488dWpVauYz8cdvx743z_La1h7zIQCsEEgu7o,14
+tensorcircuit_nightly-1.4.0.dev20251128.dist-info/RECORD,,