Repository: quantumlib/Cirq

Issue:
cirq.kraus seems to fail on PhasedXZGate
**Description of the issue**
I get a TypeError when using cirq.kraus on this gate.
**How to reproduce the issue**

```python
import cirq

op = cirq.PhasedXZGate(axis_phase_exponent=-0.5,
                           x_exponent=-0.5,
                           z_exponent=1).on(cirq.GridQubit(3, 3))
cirq.kraus(op)
```
Produces:
```
Traceback (most recent call last):
  File "/usr/local/google/home/dkafri/.config/JetBrains/PyCharmCE2021.2/scratches/scratch_62.py", line 6, in <module>
    cirq.kraus(op)
  File "/usr/local/google/home/dkafri/.virtualenvs/pyle/lib/python3.8/site-packages/cirq/protocols/kraus.py", line 189, in kraus
    raise TypeError(
TypeError: object of type '<class 'cirq.ops.gate_operation.GateOperation'>' does have a _kraus_, _mixture_ or _unitary_ method, but it returned NotImplemented.
```



**Cirq version**
I get this error both for `0.12.0.dev20210630205305` and `0.12.0` 




Discussion:
Seems like this is an error in the `kraus` protocol; it should call `cirq.unitary` to get the unitary representation of the object in question, rather than looking for a `_unitary_` method directly, because `PhasedXZGate` does work with the unitary protocol:
```
In [1]: import cirq
   ...: 
   ...: op = cirq.PhasedXZGate(axis_phase_exponent=-0.5,
   ...:                            x_exponent=-0.5,
   ...:                            z_exponent=1).on(cirq.GridQubit(3, 3))
   ...: cirq.unitary(op)
Out[1]: 
array([[ 0.5-0.5j, -0.5+0.5j],
       [-0.5+0.5j, -0.5+0.5j]])
```
We could also implement `_unitary_` directly on `PhasedXZGate`; currently it has to decompose the gate to get the unitary.
I see the same issue in mixtures as well
```python
In [1]: import cirq
   ...: op = cirq.PhasedXZGate(axis_phase_exponent=-0.5,
   ...:                            x_exponent=-0.5,
   ...:                            z_exponent=1).on(cirq.GridQubit(3, 3))
   ...: print(cirq.has_mixture(op))
   ...: cirq.mixture(op)
True
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-1-fe586f2b45b8> in <module>
      4                            z_exponent=1).on(cirq.GridQubit(3, 3))
      5 print(cirq.has_mixture(op))
----> 6 cirq.mixture(op)

~/opt/anaconda3/envs/cirq/lib/python3.8/site-packages/cirq/protocols/mixture_protocol.py in mixture(val, default)
     99         raise TypeError(f"object of type '{type(val)}' has no _mixture_ or _unitary_ method.")
    100
--> 101     raise TypeError(
    102         "object of type '{}' does have a _mixture_ or _unitary_ "
    103         "method, but it returned NotImplemented.".format(type(val))

TypeError: object of type '<class 'cirq.ops.gate_operation.GateOperation'>' does have a _mixture_ or _unitary_ method, but it returned NotImplemented.
```

The problem is that `has_[kraus|mixture]` methods check only for `_has_[kraus|mixture]_`  member functions and then fallback on `kraus/mixture` methods to check for `_kraus_ | _mixture_`. However, the former has a `allow_decompose` argument but the latter doesn't. 

I can send a PR to fix this but I have a question -- do we need a `allow_decompose` argument in `kraus`, `mixture` protocols? `unitary` protocol doesn't have the `allow_decompose` argument even though `has_unitary` does, but they both work fine because `has_unitary` doesn't depend on `unitary` as a fallback. 
Note that simply using `cirq.unitary` instead of checking for `_unitary_` is not enough. In `has_kraus`, we check whether `all(has_kraus(val) for val in decomposed_operations)` then the primary operation also `has_kraus`. Therefore, in the `kraus` method, we will need to recursively compute the `kraus` of all decomposed operations and combine their `kraus` operators by serial concatenation. 

This will be useful, for eg, to automatically compute the kraus operators of a `PauliMeasurementGate` (requested in #4127) which decomposes into pauli operations + measurement gate -- both of which have different types of kraus operators (unitary vs measurement). 
This is a duplicate of #3965, which came to a similar conclusion (add `allow_decompose` to `kraus`/`mixture`). It's important to note that limits must be placed on this decomposition, as number of Kraus operators / mixture elements can exponentially explode.
This issue is stale because it has been open 30 days with no activity. Remove stale label or comment or this will be closed in 30 days
This issue is stale because it has been open 30 days with no activity. Remove stale label or comment or this will be closed in 30 days

Resolving patch:
diff --git a/cirq-core/cirq/ops/phased_x_z_gate.py b/cirq-core/cirq/ops/phased_x_z_gate.py
--- a/cirq-core/cirq/ops/phased_x_z_gate.py
+++ b/cirq-core/cirq/ops/phased_x_z_gate.py
@@ -141,6 +141,15 @@ def _qasm_(self, args: 'cirq.QasmArgs', qubits: Tuple['cirq.Qid', ...]) -> Optio
     def _has_unitary_(self) -> bool:
         return not self._is_parameterized_()
 
+    def _unitary_(self) -> Optional[np.ndarray]:
+        """See `cirq.SupportsUnitary`."""
+        if self._is_parameterized_():
+            return None
+        z_pre = protocols.unitary(ops.Z ** -self._axis_phase_exponent)
+        x = protocols.unitary(ops.X ** self._x_exponent)
+        z_post = protocols.unitary(ops.Z ** (self._axis_phase_exponent + self._z_exponent))
+        return z_post @ x @ z_pre
+
     def _decompose_(self, qubits: Sequence['cirq.Qid']) -> 'cirq.OP_TREE':
         q = qubits[0]
         yield ops.Z(q) ** -self._axis_phase_exponent


Test patch:
diff --git a/cirq-core/cirq/ops/phased_x_z_gate_test.py b/cirq-core/cirq/ops/phased_x_z_gate_test.py
--- a/cirq-core/cirq/ops/phased_x_z_gate_test.py
+++ b/cirq-core/cirq/ops/phased_x_z_gate_test.py
@@ -199,6 +199,22 @@ def test_from_matrix_close_unitary(unitary: np.ndarray):
     )
 
 
+@pytest.mark.parametrize(
+    'unitary',
+    [
+        cirq.testing.random_unitary(2),
+        cirq.testing.random_unitary(2),
+        cirq.testing.random_unitary(2),
+        np.array([[0, 1], [1j, 0]]),
+    ],
+)
+def test_from_matrix_close_kraus(unitary: np.ndarray):
+    gate = cirq.PhasedXZGate.from_matrix(unitary)
+    kraus = cirq.kraus(gate)
+    assert len(kraus) == 1
+    cirq.testing.assert_allclose_up_to_global_phase(kraus[0], unitary, atol=1e-8)
+
+
 def test_protocols():
     a = random.random()
     b = random.random()