fix to unitary_kraus

Author: refraction-ray

CHANGELOG.md

@@ -26,6 +26,12 @@
 
 - Add ccx as another alias for toffoli gate
 
+### Changed
+
+- Seperate channel auto register for circuit class with unitary and general case
+
+- The old standlone depolarizing implementation can now be called via `c.depolarizing_reference`
+
 ### Fixed
 
 - Move `iswap` gate to vgates list
@@ -36,6 +42,8 @@
 
 - Fix the use of rev over rev jacobian for hessian method and back to the efficient solution of fwd over rev order due to the solution of nested vmap issue on tf backend
 
+- Identify potential bug in `unitary_kraus2` implementation, change to `unitary_kraus` instead
+
 ## 0.4.0
 
 ### Added

tensorcircuit/applications/layers.py

@@ -213,8 +213,8 @@ def f(
                 qubit2,
                 -symbol0 * g[e[0]][e[1]].get("weight", 1.0) * 2,
             )  ## should be better as * 2 # e^{-i\theta H}, H=-ZZ
-            circuit.depolarizing(e[0], px=params[0], py=params[1], pz=params[2])
-            circuit.depolarizing(e[1], px=params[0], py=params[1], pz=params[2])
+            circuit.depolarizing(e[0], px=params[0], py=params[1], pz=params[2])  # type: ignore
+            circuit.depolarizing(e[1], px=params[0], py=params[1], pz=params[2])  # type: ignore
         return circuit
 
     f.__doc__ = """%slayer_bitflip_mc""" % gates
@@ -238,8 +238,8 @@ def f(
                 -symbol[i] * g[e[0]][e[1]].get("weight", 1.0) * 2,  # type: ignore
             )
             ## should be better as * 2 # e^{-i\theta H}, H=-ZZ
-            circuit.depolarizing(e[0], px=params[0], py=params[1], pz=params[2])
-            circuit.depolarizing(e[1], px=params[0], py=params[1], pz=params[2])
+            circuit.depolarizing(e[0], px=params[0], py=params[1], pz=params[2])  # type: ignore
+            circuit.depolarizing(e[1], px=params[0], py=params[1], pz=params[2])  # type: ignore
         return circuit
 
     f.__doc__ = """any%slayer_bitflip_mc""" % gates
@@ -288,8 +288,8 @@ def anyswaplayer_bitflip_mc(
             unitary=array_to_tensor(_swap_matrix),
             theta=symbol[i] * g[e[0]][e[1]].get("weight", 1.0),
         )
-        circuit.depolarizing(e[0], px=px, py=py, pz=pz)
-        circuit.depolarizing(e[1], px=px, py=py, pz=pz)
+        circuit.depolarizing(e[0], px=px, py=py, pz=pz)  # type: ignore
+        circuit.depolarizing(e[1], px=px, py=py, pz=pz)  # type: ignore
     return circuit
 
 
@@ -324,7 +324,7 @@ def bitfliplayer(ci: DMCircuit, g: Graph, px: float, py: float, pz: float) -> No
 def bitfliplayer_mc(ci: Circuit, g: Graph, px: float, py: float, pz: float) -> None:
     n = len(g.nodes)
     for i in range(n):
-        ci.depolarizing(i, px=px, py=py, pz=pz)
+        ci.depolarizing(i, px=px, py=py, pz=pz)  # type: ignore
     bitfliplayer.__repr__ = """bitfliplayer_mc"""  # type: ignore
     bitfliplayer.__trainable__ = True  # type: ignore
 

tensorcircuit/circuit.py

@@ -255,7 +255,7 @@ def depolarizing2(
         # building time and running time are similar
 
     # overwritten now, deprecated
-    def depolarizing(
+    def depolarizing_reference(
         self,
         index: int,
         *,
@@ -321,6 +321,8 @@ def unitary_kraus2(
         status: Optional[float] = None,
         name: Optional[str] = None,
     ) -> 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
         # still jittable
         # speed is similar to ``unitary_kraus``
@@ -545,7 +547,7 @@ def calculate_kraus_p(i: int) -> Tensor:
             for w, k in zip(prob, kraus_tensor)
         ]
 
-        return self.unitary_kraus2(
+        return self.unitary_kraus(
             new_kraus, *index, prob=prob, status=status, name=name
         )
 
@@ -579,7 +581,7 @@ def general_kraus(
 
     @staticmethod
     def apply_general_kraus_delayed(
-        krausf: Callable[..., Sequence[Gate]]
+        krausf: Callable[..., Sequence[Gate]], is_unitary: bool = False
     ) -> Callable[..., None]:
         def apply(
             self: "Circuit",
@@ -589,17 +591,26 @@ def apply(
             **vars: float,
         ) -> None:
             kraus = krausf(**vars)
-            self.apply_general_kraus(kraus, *index, status=status, name=name)
+            if not is_unitary:
+                self.apply_general_kraus(kraus, *index, status=status, name=name)
+            else:
+                self.unitary_kraus(kraus, *index, status=status, name=name)
 
         return apply
 
     @classmethod
     def _meta_apply_channels(cls) -> None:
         for k in channels.channels:
+            if k in ["depolarizing", "generaldepolarizing"]:
+                is_unitary = True
+            else:
+                is_unitary = False
             setattr(
                 cls,
                 k,
-                cls.apply_general_kraus_delayed(getattr(channels, k + "channel")),
+                cls.apply_general_kraus_delayed(
+                    getattr(channels, k + "channel"), is_unitary=is_unitary
+                ),
             )
             doc = """
             Apply %s quantum channel on the circuit.