add docstrings;fix jax to_dlpack;add with_prob impl for stabilizercircuit.measure

Author: refraction-ray

CHANGELOG.md

@@ -10,10 +10,14 @@
 
 - Add `circuit.amplitude_before()` method to return the corresponding tensornetwork nodes.
 
+- Add `with_prob` for `stabilizercircuit.measure()`.
+
 ### 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/backends/jax_backend.py

@@ -442,7 +442,14 @@ def from_dlpack(self, a: Any) -> Tensor:
     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 @@ def depolarizing2(
         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 @@ def unitary_kraus2(
         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