make everything easier.

Author: WeiguoMa

tensorcircuit/circuit.py

@@ -391,7 +391,6 @@ def unitary_kraus(
         prob: Optional[Sequence[float]] = None,
         status: Optional[float] = None,
         name: Optional[str] = None,
-        return_gate: bool = False,
     ) -> Tensor:
         """
         Apply unitary gates in ``kraus`` randomly based on corresponding ``prob``.
@@ -423,7 +422,6 @@ def index2gate(r: Tensor, kraus: Sequence[Tensor]) -> Tensor:
             status=status,
             get_gate_from_index=index2gate,
             name=name,
-            return_gate=return_gate,
         )
 
     def _unitary_kraus_template(
@@ -436,7 +434,6 @@ def _unitary_kraus_template(
             Callable[[Tensor, Sequence[Tensor]], Tensor]
         ] = None,
         name: Optional[str] = None,
-        return_gate: bool = False,
     ) -> Tensor:  # DRY
         sites = len(index)
         kraus = [k.tensor if isinstance(k, tn.Node) else k for k in kraus]
@@ -478,9 +475,7 @@ def step_function(x: Tensor) -> Tensor:
             raise ValueError("no `get_gate_from_index` implementation is provided")
         g = get_gate_from_index(r, kraus)
         g = backend.reshape(g, [self._d for _ in range(sites * 2)])
-        if return_gate:
-            return r, g
-        self.any(*index, unitary=g, name=name)  # type: ignore
+        self.any(*index, unitary=g, name=name, dim=self._d)  # type: ignore
         return r
 
     def _general_kraus_tf(
@@ -558,7 +553,6 @@ def _general_kraus_2(
         status: Optional[float] = None,
         with_prob: bool = False,
         name: Optional[str] = None,
-        return_gate: bool = False,
     ) -> Tensor:
         # the graph building time is frustratingly slow, several minutes
         # though running time is in terms of ms
@@ -611,7 +605,6 @@ def calculate_kraus_p(i: int) -> Tensor:
             prob=prob,
             status=status,
             name=name,
-            return_gate=return_gate,
         )
         if not with_prob:
             return pick
@@ -625,7 +618,6 @@ def general_kraus(
         status: Optional[float] = None,
         with_prob: bool = False,
         name: Optional[str] = None,
-        return_gate: bool = False,
     ) -> Tensor:
         """
         Monte Carlo trajectory simulation of general Kraus channel whose Kraus operators cannot be
@@ -650,7 +642,6 @@ def general_kraus(
             status=status,
             with_prob=with_prob,
             name=name,
-            return_gate=return_gate,
         )
 
     apply_general_kraus = general_kraus

tensorcircuit/quditcircuit.py

@@ -695,24 +695,14 @@ def general_kraus(
             when the random number will be generated automatically
         :type status: Optional[float], optional
         """
-        res = self._circ.general_kraus(
+        return self._circ.general_kraus(
             kraus,
             *index,
             status=status,
             with_prob=with_prob,
             name=name,
-            return_gate=True,
         )
 
-        if with_prob:
-            (pick, gate), prob = res
-            self.any(*index, unitary=gate, name=name)
-            return pick, prob
-        else:
-            pick, gate = res
-            self.any(*index, unitary=gate, name=name)
-            return pick
-
     def unitary_kraus(
         self,
         kraus: Sequence[Gate],
@@ -734,13 +724,10 @@ def unitary_kraus(
         :return: shape [] int dtype tensor indicates which kraus gate is actually applied
         :rtype: Tensor
         """
-        r, g = self._circ.unitary_kraus(
+        return self._circ.unitary_kraus(
             kraus,
             *index,
             prob=prob,
             status=status,
             name=name,
-            return_gate=True,
         )
-        self.any(*index, unitary=g, name=name)
-        return r