Repository: quantumlib/Cirq

Issue:
Public PlusGate
**Is your feature request related to a use case or problem? Please describe.**

We have `PlusGate` class implementation duplicated several times in tests, and there's now a copy of that class used in the internal implementation of `ResetChannel`.


**Describe the solution you'd like**

Should we just make a public implementation? (This could be a good first issue, no particular QC knowledge needed)


**[optional] Describe alternatives/workarounds you've considered**

Not make it public.


**What is the urgency from your perspective for this issue? Is it blocking important work?**
<!-- Please choose one and remove the others -->

P3 - I'm not really blocked by it, it is an idea I'd like to discuss / suggestion based on principle 

<!-- [optional] additional comment / context -->


Discussion:
In light of https://github.com/quantumlib/Cirq/issues/3190 (if I'm understanding correctly), would another option be to rewrite `XPowGate` as a `PlusGate`, with optional `dimension` and `increment` arguments? 

Assuming my understanding of the issue is correct, this approach would match what we do with `IdentityGate`, where `cirq.I` is a single-qubit gate by default, but you can initialize `cirq.IdentityGate(dimension=5)` if you need to.
The PlusGate as implemented is probably insufficient. We'd probably want to figure out the EigenGate representation of it before adding it to our public gates. So one consideration is how it works under non-integer exponents. Integer exponents would just multiply the `increment` (meaning we don't actually need an `increment`, and can just use the exponent). However non-integer exponents, I'm not sure how that would work. There's an `_eigen_components` function that, if implemented, might allow things to "just work". I *assume* PlusGate has well-defined Eigen components, but not sure.
@Strilanc @viathor I'm having trouble understanding what `_eigen_components` does and how/whether it would work for PlusGate. Any google results I see just talk about the unitaries. Do you have any references?
I think I figured the math out, but it's coming out to there being no eigen components that satisfy a PlusGate. Eigenvalues (for a 3-value PlusGate) are the third roots of `i`, so half-turns are `[0, 2/3, 4/3]`. But the symmetry about the real line implies that the 2nd eigenspace has to equal the 3rd, so the complex values can cancel out. But if they're equal then, well, there's nothing to differentiate them. But they have to be different. So is a PlusGate not able to be an Eigen gate?

cc @cduck 

Note I haven't done linear algebra in 28 years (and even that was a summer crash course), so.
I figured out the math, reversing the logic in `EigenGate._unitary_` to get the required eigencomponents, which end up looking like
```
(0, np.array([[1/3, -2/3, 0], [0, 1/3, -2/3], [-2/3, 0, 1/3]])),
(2/3, np.array([[1/3, 1/3, 0], [0, 1/3, 1/3], [1/3, 0, 1/3]])),
(4/3, np.array([[1/3, 1/3, 0], [0, 1/3, 1/3], [1/3, 0, 1/3]])),
```
Plugging that into a new EigenGate `Plus` gives `unitary(Plus)`
```
[[0.+0.j 1.+0.j 0.+0.j]
 [0.+0.j 0.+0.j 1.+0.j]
 [1.+0.j 0.+0.j 0.+0.j]]
```
However `unitary(Plus**0.5)` gives:
```
[[0.+0.67j 0.-0.33j 0.+0.j  ]
 [0.+0.j   0.+0.67j 0.-0.33j]
 [0.-0.33j 0.+0.j   0.+0.67j]]
```
Looks like something that could plausibly be a unitary, but it isn't. Multiplying by adjoint gives:
```
[[ 0.56+0.j -0.22+0.j -0.22+0.j]
 [-0.22+0.j  0.56+0.j -0.22+0.j]
 [-0.22+0.j -0.22+0.j  0.56+0.j]]
```

So, I guess a PlusGate isn't an EigenGate (at least in dimension 3)?
I tried again, trying some linear algebra directly instead of reversing the existing `EigenGate._unitary_`. (https://en.wikipedia.org/wiki/Square_root_of_a_matrix#Matrices_with_distinct_eigenvalues)

Using that I ended up with what the unitary for `Plus**0,5` *should* be.
```
[[0.33+0.58j 0.33-0.58j 0.33-0.j  ]
 [0.33-0.j   0.33+0.58j 0.33-0.58j]
 [0.33-0.58j 0.33-0.j   0.33+0.58j]]
```
Indeed when I plug that into a `MatrixGate`, and simulate it, it provides the following fun moment steps:
```
[1. 0. 0.]
[0.67 0.33 0.67]
[0. 0. 1.]
[0.33 0.67 0.67]
[0. 1. 0.]
[0.67 0.67 0.33]
[1. 0. 0.]
```
However I can't figure out what the `_eigen_components` need to be in order to make `EigenGate._unitary_` output that value when `exponent=0.5`. (In fact, it seems that there are no such components, as when I worked the algorithm out backwards in my previous comment, that ended up being the *unique* solution that satisfied the `exponent=0` and `exponent=1` unitaries for what `Plus` should be.

So, it seems that `Plus` *is* an EigenGate. So is the implementation of `EigenGate._unitary_` wrong since it seems not to be able to represent this case? (But it works for so many other cases, so I have a hard time believing that it would be).
Okay, looks like these components are projection matrices onto the eigenvectors. So, pretty sure I see the pattern for making the eigenvectors for a PlusGate (roots of unity and their powers). Still need to look up the formula for projecting onto a complex vector, but thinking this is getting close. 
That seems to have worked for dimension 3. If I change X to:

```python
    def _qid_shape_(self) -> Tuple[int, ...]:
        return (3,)

    def _eigen_components(self) -> List[Tuple[float, np.ndarray]]:
        L = (-1+np.sqrt(3)*1j)/2  # 3rd root of unity
        L2 = L ** 2
        return [
            (0, np.array([
                [1, 1, 1],
                [1, 1, 1],
                [1, 1, 1],
            ]) / 3),
            (2/3, np.array([
                [1, L, L2],
                [L2, 1, L],
                [L, L2, 1],
            ]) / 3),
            (4/3, np.array([
                [1, L2, L],
                [L, 1, L2],
                [L2, L, 1],
            ]) / 3),
            # (0, np.array([[0.5, 0.5], [0.5, 0.5]])),
            # (1, np.array([[0.5, -0.5], [-0.5, 0.5]])),
        ]
```
then it provides the nice step function, and `X**0,5` moment steps are
```
[[1.  +0.j   0.  +0.j   0.  +0.j  ]
 [0.33+0.58j 0.33-0.58j 0.33-0.j  ]
 [0.  +0.j   1.  -0.j   0.  +0.j  ]
 [0.33+0.j   0.33+0.58j 0.33-0.58j]
 [0.  +0.j   0.  +0.j   1.  -0.j  ]
 [0.33-0.58j 0.33-0.j   0.33+0.58j]
 [1.  +0.j   0.  +0.j   0.  +0.j  ]]
```

So the answer to why QM requires complex numbers is: so that we can count to 3 by halves. :)

Anyway that seems to be a fairly straightforward pattern. I'll see what I can get for higher dimensions (I'm guessing it may be slightly different for non-primes) and then create a PR. (Not sure what I did wrong in my previous calculation that suggested such a decomposition wasn't possible, but oh well).

Resolving patch:
diff --git a/cirq-core/cirq/ops/common_channels.py b/cirq-core/cirq/ops/common_channels.py
--- a/cirq-core/cirq/ops/common_channels.py
+++ b/cirq-core/cirq/ops/common_channels.py
@@ -717,23 +717,6 @@ def _act_on_(self, sim_state: 'cirq.SimulationStateBase', qubits: Sequence['cirq
         if len(qubits) != 1:
             return NotImplemented
 
-        class PlusGate(raw_types.Gate):
-            """A qudit gate that increments a qudit state mod its dimension."""
-
-            def __init__(self, dimension, increment=1):
-                self.dimension = dimension
-                self.increment = increment % dimension
-
-            def _qid_shape_(self):
-                return (self.dimension,)
-
-            def _unitary_(self):
-                inc = (self.increment - 1) % self.dimension + 1
-                u = np.empty((self.dimension, self.dimension))
-                u[inc:] = np.eye(self.dimension)[:-inc]
-                u[:inc] = np.eye(self.dimension)[-inc:]
-                return u
-
         from cirq.sim import simulation_state
 
         if (
@@ -741,7 +724,7 @@ def _unitary_(self):
             and not sim_state.can_represent_mixed_states
         ):
             result = sim_state._perform_measurement(qubits)[0]
-            gate = PlusGate(self.dimension, self.dimension - result)
+            gate = common_gates.XPowGate(dimension=self.dimension) ** (self.dimension - result)
             protocols.act_on(gate, sim_state, qubits)
             return True
 
diff --git a/cirq-core/cirq/ops/common_gates.py b/cirq-core/cirq/ops/common_gates.py
--- a/cirq-core/cirq/ops/common_gates.py
+++ b/cirq-core/cirq/ops/common_gates.py
@@ -87,11 +87,19 @@ class XPowGate(eigen_gate.EigenGate):
     `cirq.X`, the Pauli X gate, is an instance of this gate at `exponent=1`.
     """
 
+    _eigencomponents: Dict[int, List[Tuple[float, np.ndarray]]] = {}
+
+    def __init__(
+        self, *, exponent: value.TParamVal = 1.0, global_shift: float = 0.0, dimension: int = 2
+    ):
+        super().__init__(exponent=exponent, global_shift=global_shift)
+        self._dimension = dimension
+
     def _num_qubits_(self) -> int:
         return 1
 
     def _apply_unitary_(self, args: 'protocols.ApplyUnitaryArgs') -> Optional[np.ndarray]:
-        if self._exponent != 1:
+        if self._exponent != 1 or self._dimension != 2:
             return NotImplemented
         zero = args.subspace_index(0)
         one = args.subspace_index(1)
@@ -108,10 +116,27 @@ def in_su2(self) -> 'Rx':
 
     def with_canonical_global_phase(self) -> 'XPowGate':
         """Returns an equal-up-global-phase standardized form of the gate."""
-        return XPowGate(exponent=self._exponent)
+        return XPowGate(exponent=self._exponent, dimension=self._dimension)
+
+    def _qid_shape_(self) -> Tuple[int, ...]:
+        return (self._dimension,)
 
     def _eigen_components(self) -> List[Tuple[float, np.ndarray]]:
-        return [(0, np.array([[0.5, 0.5], [0.5, 0.5]])), (1, np.array([[0.5, -0.5], [-0.5, 0.5]]))]
+        if self._dimension not in XPowGate._eigencomponents:
+            components = []
+            root = 1j ** (4 / self._dimension)
+            for i in range(self._dimension):
+                half_turns = i * 2 / self._dimension
+                v = np.array([root ** (i * j) / self._dimension for j in range(self._dimension)])
+                m = np.array([np.roll(v, j) for j in range(self._dimension)])
+                components.append((half_turns, m))
+            XPowGate._eigencomponents[self._dimension] = components
+        return XPowGate._eigencomponents[self._dimension]
+
+    def _with_exponent(self, exponent: 'cirq.TParamVal') -> 'cirq.XPowGate':
+        return XPowGate(
+            exponent=exponent, global_shift=self._global_shift, dimension=self._dimension
+        )
 
     def _decompose_into_clifford_with_qubits_(self, qubits):
         from cirq.ops.clifford_gate import SingleQubitCliffordGate
@@ -127,7 +152,7 @@ def _decompose_into_clifford_with_qubits_(self, qubits):
         return NotImplemented
 
     def _trace_distance_bound_(self) -> Optional[float]:
-        if self._is_parameterized_():
+        if self._is_parameterized_() or self._dimension != 2:
             return None
         return abs(np.sin(self._exponent * 0.5 * np.pi))
 
@@ -179,7 +204,7 @@ def controlled(
         return result
 
     def _pauli_expansion_(self) -> value.LinearDict[str]:
-        if protocols.is_parameterized(self):
+        if protocols.is_parameterized(self) or self._dimension != 2:
             return NotImplemented
         phase = 1j ** (2 * self._exponent * (self._global_shift + 0.5))
         angle = np.pi * self._exponent / 2
@@ -220,7 +245,7 @@ def _phase_by_(self, phase_turns, qubit_index):
         )
 
     def _has_stabilizer_effect_(self) -> Optional[bool]:
-        if self._is_parameterized_():
+        if self._is_parameterized_() or self._dimension != 2:
             return None
         return self.exponent % 0.5 == 0
 
@@ -232,13 +257,19 @@ def __str__(self) -> str:
         return f'XPowGate(exponent={self._exponent}, global_shift={self._global_shift!r})'
 
     def __repr__(self) -> str:
-        if self._global_shift == 0:
+        if self._global_shift == 0 and self._dimension == 2:
             if self._exponent == 1:
                 return 'cirq.X'
             return f'(cirq.X**{proper_repr(self._exponent)})'
-        return 'cirq.XPowGate(exponent={}, global_shift={!r})'.format(
-            proper_repr(self._exponent), self._global_shift
-        )
+        args = []
+        if self._exponent != 1:
+            args.append(f'exponent={proper_repr(self._exponent)}')
+        if self._global_shift != 0:
+            args.append(f'global_shift={self._global_shift}')
+        if self._dimension != 2:
+            args.append(f'dimension={self._dimension}')
+        all_args = ', '.join(args)
+        return f'cirq.XPowGate({all_args})'
 
 
 class Rx(XPowGate):
@@ -478,6 +509,14 @@ class ZPowGate(eigen_gate.EigenGate):
     `cirq.Z`, the Pauli Z gate, is an instance of this gate at `exponent=1`.
     """
 
+    _eigencomponents: Dict[int, List[Tuple[float, np.ndarray]]] = {}
+
+    def __init__(
+        self, *, exponent: value.TParamVal = 1.0, global_shift: float = 0.0, dimension: int = 2
+    ):
+        super().__init__(exponent=exponent, global_shift=global_shift)
+        self._dimension = dimension
+
     def _num_qubits_(self) -> int:
         return 1
 
@@ -485,9 +524,10 @@ def _apply_unitary_(self, args: 'protocols.ApplyUnitaryArgs') -> Optional[np.nda
         if protocols.is_parameterized(self):
             return None
 
-        one = args.subspace_index(1)
-        c = 1j ** (self._exponent * 2)
-        args.target_tensor[one] *= c
+        for i in range(1, self._dimension):
+            subspace = args.subspace_index(i)
+            c = 1j ** (self._exponent * 4 * i / self._dimension)
+            args.target_tensor[subspace] *= c
         p = 1j ** (2 * self._exponent * self._global_shift)
         if p != 1:
             args.target_tensor *= p
@@ -512,7 +552,7 @@ def in_su2(self) -> 'Rz':
 
     def with_canonical_global_phase(self) -> 'ZPowGate':
         """Returns an equal-up-global-phase standardized form of the gate."""
-        return ZPowGate(exponent=self._exponent)
+        return ZPowGate(exponent=self._exponent, dimension=self._dimension)
 
     def controlled(
         self,
@@ -561,16 +601,32 @@ def controlled(
             )
         return result
 
+    def _qid_shape_(self) -> Tuple[int, ...]:
+        return (self._dimension,)
+
     def _eigen_components(self) -> List[Tuple[float, np.ndarray]]:
-        return [(0, np.diag([1, 0])), (1, np.diag([0, 1]))]
+        if self._dimension not in ZPowGate._eigencomponents:
+            components = []
+            for i in range(self._dimension):
+                half_turns = i * 2 / self._dimension
+                m = np.zeros((self._dimension, self._dimension))
+                m[i][i] = 1
+                components.append((half_turns, m))
+            ZPowGate._eigencomponents[self._dimension] = components
+        return ZPowGate._eigencomponents[self._dimension]
+
+    def _with_exponent(self, exponent: 'cirq.TParamVal') -> 'cirq.ZPowGate':
+        return ZPowGate(
+            exponent=exponent, global_shift=self._global_shift, dimension=self._dimension
+        )
 
     def _trace_distance_bound_(self) -> Optional[float]:
-        if self._is_parameterized_():
+        if self._is_parameterized_() or self._dimension != 2:
             return None
         return abs(np.sin(self._exponent * 0.5 * np.pi))
 
     def _pauli_expansion_(self) -> value.LinearDict[str]:
-        if protocols.is_parameterized(self):
+        if protocols.is_parameterized(self) or self._dimension != 2:
             return NotImplemented
         phase = 1j ** (2 * self._exponent * (self._global_shift + 0.5))
         angle = np.pi * self._exponent / 2
@@ -580,7 +636,7 @@ def _phase_by_(self, phase_turns: float, qubit_index: int):
         return self
 
     def _has_stabilizer_effect_(self) -> Optional[bool]:
-        if self._is_parameterized_():
+        if self._is_parameterized_() or self._dimension != 2:
             return None
         return self.exponent % 0.5 == 0
 
@@ -630,7 +686,7 @@ def __str__(self) -> str:
         return f'ZPowGate(exponent={self._exponent}, global_shift={self._global_shift!r})'
 
     def __repr__(self) -> str:
-        if self._global_shift == 0:
+        if self._global_shift == 0 and self._dimension == 2:
             if self._exponent == 0.25:
                 return 'cirq.T'
             if self._exponent == -0.25:
@@ -642,9 +698,15 @@ def __repr__(self) -> str:
             if self._exponent == 1:
                 return 'cirq.Z'
             return f'(cirq.Z**{proper_repr(self._exponent)})'
-        return 'cirq.ZPowGate(exponent={}, global_shift={!r})'.format(
-            proper_repr(self._exponent), self._global_shift
-        )
+        args = []
+        if self._exponent != 1:
+            args.append(f'exponent={proper_repr(self._exponent)}')
+        if self._global_shift != 0:
+            args.append(f'global_shift={self._global_shift}')
+        if self._dimension != 2:
+            args.append(f'dimension={self._dimension}')
+        all_args = ', '.join(args)
+        return f'cirq.ZPowGate({all_args})'
 
     def _commutes_on_qids_(
         self, qids: 'Sequence[cirq.Qid]', other: Any, *, atol: float = 1e-8


Test patch:
diff --git a/cirq-core/cirq/ops/common_gates_test.py b/cirq-core/cirq/ops/common_gates_test.py
--- a/cirq-core/cirq/ops/common_gates_test.py
+++ b/cirq-core/cirq/ops/common_gates_test.py
@@ -1116,3 +1116,190 @@ def test_approx_eq():
     assert not cirq.approx_eq(cirq.Y**0.1, cirq.Y**0.2, atol=0.05)
     assert cirq.approx_eq(cirq.X**0.1, cirq.X**0.2, atol=0.3)
     assert not cirq.approx_eq(cirq.X**0.1, cirq.X**0.2, atol=0.05)
+
+
+def test_xpow_dim_3():
+    x = cirq.XPowGate(dimension=3)
+    # fmt: off
+    expected = [
+        [0, 0, 1],
+        [1, 0, 0],
+        [0, 1, 0],
+    ]
+    # fmt: on
+    assert np.allclose(cirq.unitary(x), expected)
+
+    sim = cirq.Simulator()
+    circuit = cirq.Circuit([x(cirq.LineQid(0, 3)) ** 0.5] * 6)
+    svs = [step.state_vector() for step in sim.simulate_moment_steps(circuit)]
+    # fmt: off
+    expected = [
+        [0.67, 0.67, 0.33],
+        [0.0, 1.0, 0.0],
+        [0.33, 0.67, 0.67],
+        [0.0, 0.0, 1.0],
+        [0.67, 0.33, 0.67],
+        [1.0, 0.0, 0.0],
+    ]
+    # fmt: on
+    assert np.allclose(np.abs(svs), expected, atol=1e-2)
+
+
+def test_xpow_dim_4():
+    x = cirq.XPowGate(dimension=4)
+    # fmt: off
+    expected = [
+        [0, 0, 0, 1],
+        [1, 0, 0, 0],
+        [0, 1, 0, 0],
+        [0, 0, 1, 0],
+    ]
+    # fmt: on
+    assert np.allclose(cirq.unitary(x), expected)
+
+    sim = cirq.Simulator()
+    circuit = cirq.Circuit([x(cirq.LineQid(0, 4)) ** 0.5] * 8)
+    svs = [step.state_vector() for step in sim.simulate_moment_steps(circuit)]
+    # fmt: off
+    expected = [
+        [0.65, 0.65, 0.27, 0.27],
+        [0.0, 1.0, 0.0, 0.0],
+        [0.27, 0.65, 0.65, 0.27],
+        [0.0, 0.0, 1.0, 0.0],
+        [0.27, 0.27, 0.65, 0.65],
+        [0.0, 0.0, 0.0, 1.0],
+        [0.65, 0.27, 0.27, 0.65],
+        [1.0, 0.0, 0.0, 0.0],
+    ]
+    # fmt: on
+    assert np.allclose(np.abs(svs), expected, atol=1e-2)
+
+
+def test_zpow_dim_3():
+    L = np.exp(2 * np.pi * 1j / 3)
+    L2 = L**2
+    z = cirq.ZPowGate(dimension=3)
+    # fmt: off
+    expected = [
+        [1, 0, 0],
+        [0, L, 0],
+        [0, 0, L2],
+    ]
+    # fmt: on
+    assert np.allclose(cirq.unitary(z), expected)
+
+    sim = cirq.Simulator()
+    circuit = cirq.Circuit([z(cirq.LineQid(0, 3)) ** 0.5] * 6)
+    svs = [step.state_vector() for step in sim.simulate_moment_steps(circuit, initial_state=0)]
+    expected = [[1, 0, 0]] * 6
+    assert np.allclose((svs), expected)
+
+    svs = [step.state_vector() for step in sim.simulate_moment_steps(circuit, initial_state=1)]
+    # fmt: off
+    expected = [
+        [0, L**0.5, 0],
+        [0, L**1.0, 0],
+        [0, L**1.5, 0],
+        [0, L**2.0, 0],
+        [0, L**2.5, 0],
+        [0, 1, 0],
+    ]
+    # fmt: on
+    assert np.allclose((svs), expected)
+
+    svs = [step.state_vector() for step in sim.simulate_moment_steps(circuit, initial_state=2)]
+    # fmt: off
+    expected = [
+        [0, 0, L],
+        [0, 0, L2],
+        [0, 0, 1],
+        [0, 0, L],
+        [0, 0, L2],
+        [0, 0, 1],
+    ]
+    # fmt: on
+    assert np.allclose((svs), expected)
+
+
+def test_zpow_dim_4():
+    z = cirq.ZPowGate(dimension=4)
+    # fmt: off
+    expected = [
+        [1, 0, 0, 0],
+        [0, 1j, 0, 0],
+        [0, 0, -1, 0],
+        [0, 0, 0, -1j],
+    ]
+    # fmt: on
+    assert np.allclose(cirq.unitary(z), expected)
+
+    sim = cirq.Simulator()
+    circuit = cirq.Circuit([z(cirq.LineQid(0, 4)) ** 0.5] * 8)
+    svs = [step.state_vector() for step in sim.simulate_moment_steps(circuit, initial_state=0)]
+    expected = [[1, 0, 0, 0]] * 8
+    assert np.allclose((svs), expected)
+
+    svs = [step.state_vector() for step in sim.simulate_moment_steps(circuit, initial_state=1)]
+    # fmt: off
+    expected = [
+        [0, 1j**0.5, 0, 0],
+        [0, 1j**1.0, 0, 0],
+        [0, 1j**1.5, 0, 0],
+        [0, 1j**2.0, 0, 0],
+        [0, 1j**2.5, 0, 0],
+        [0, 1j**3.0, 0, 0],
+        [0, 1j**3.5, 0, 0],
+        [0, 1, 0, 0],
+    ]
+    # fmt: on
+    assert np.allclose(svs, expected)
+
+    svs = [step.state_vector() for step in sim.simulate_moment_steps(circuit, initial_state=2)]
+    # fmt: off
+    expected = [
+        [0, 0, 1j, 0],
+        [0, 0, -1, 0],
+        [0, 0, -1j, 0],
+        [0, 0, 1, 0],
+        [0, 0, 1j, 0],
+        [0, 0, -1, 0],
+        [0, 0, -1j, 0],
+        [0, 0, 1, 0],
+    ]
+    # fmt: on
+    assert np.allclose(svs, expected)
+
+    svs = [step.state_vector() for step in sim.simulate_moment_steps(circuit, initial_state=3)]
+    # fmt: off
+    expected = [
+        [0, 0, 0, 1j**1.5],
+        [0, 0, 0, 1j**3],
+        [0, 0, 0, 1j**0.5],
+        [0, 0, 0, 1j**2],
+        [0, 0, 0, 1j**3.5],
+        [0, 0, 0, 1j**1],
+        [0, 0, 0, 1j**2.5],
+        [0, 0, 0, 1],
+    ]
+    # fmt: on
+    assert np.allclose(svs, expected)
+
+
+def test_wrong_dims():
+    x3 = cirq.XPowGate(dimension=3)
+    with pytest.raises(ValueError, match='Wrong shape'):
+        _ = x3.on(cirq.LineQubit(0))
+    with pytest.raises(ValueError, match='Wrong shape'):
+        _ = x3.on(cirq.LineQid(0, dimension=4))
+
+    z3 = cirq.ZPowGate(dimension=3)
+    with pytest.raises(ValueError, match='Wrong shape'):
+        _ = z3.on(cirq.LineQubit(0))
+    with pytest.raises(ValueError, match='Wrong shape'):
+        _ = z3.on(cirq.LineQid(0, dimension=4))
+
+    with pytest.raises(ValueError, match='Wrong shape'):
+        _ = cirq.X.on(cirq.LineQid(0, dimension=3))
+
+    with pytest.raises(ValueError, match='Wrong shape'):
+        _ = cirq.Z.on(cirq.LineQid(0, dimension=3))
diff --git a/cirq-core/cirq/sim/density_matrix_simulator_test.py b/cirq-core/cirq/sim/density_matrix_simulator_test.py
--- a/cirq-core/cirq/sim/density_matrix_simulator_test.py
+++ b/cirq-core/cirq/sim/density_matrix_simulator_test.py
@@ -24,24 +24,6 @@
 import cirq.testing
 
 
-class PlusGate(cirq.Gate):
-    """A qudit gate that increments a qudit state mod its dimension."""
-
-    def __init__(self, dimension, increment=1):
-        self.dimension = dimension
-        self.increment = increment % dimension
-
-    def _qid_shape_(self):
-        return (self.dimension,)
-
-    def _unitary_(self):
-        inc = (self.increment - 1) % self.dimension + 1
-        u = np.empty((self.dimension, self.dimension))
-        u[inc:] = np.eye(self.dimension)[:-inc]
-        u[:inc] = np.eye(self.dimension)[-inc:]
-        return u
-
-
 class _TestMixture(cirq.Gate):
     def __init__(self, gate_options):
         self.gate_options = gate_options
@@ -135,8 +117,8 @@ def test_run_qudit_increments(dtype: Type[np.number], split: bool):
     for b0 in [0, 1, 2]:
         for b1 in [0, 1, 2, 3]:
             circuit = cirq.Circuit(
-                [PlusGate(3, 1)(q0)] * b0,
-                [PlusGate(4, 1)(q1)] * b1,
+                [cirq.XPowGate(dimension=3)(q0)] * b0,
+                [cirq.XPowGate(dimension=4)(q1)] * b1,
                 cirq.measure(q0),
                 cirq.measure(q1),
             )
@@ -184,7 +166,13 @@ def test_run_mixture(dtype: Type[np.number], split: bool):
 @pytest.mark.parametrize('split', [True, False])
 def test_run_qudit_mixture(dtype: Type[np.number], split: bool):
     q0, q1 = cirq.LineQid.for_qid_shape((3, 2))
-    mixture = _TestMixture([PlusGate(3, 0), PlusGate(3, 1), PlusGate(3, 2)])
+    mixture = _TestMixture(
+        [
+            cirq.XPowGate(dimension=3) ** 0,
+            cirq.XPowGate(dimension=3),
+            cirq.XPowGate(dimension=3) ** 2,
+        ]
+    )
     circuit = cirq.Circuit(mixture(q0), cirq.measure(q0), cirq.measure(q1))
     simulator = cirq.DensityMatrixSimulator(dtype=dtype, split_untangled_states=split)
     result = simulator.run(circuit, repetitions=100)
@@ -249,7 +237,11 @@ def _kraus_(self):
 
     q0, q1 = cirq.LineQid.for_qid_shape((3, 4))
     circuit = cirq.Circuit(
-        PlusGate(3, 2)(q0), TestChannel()(q0), TestChannel()(q0), cirq.measure(q0), cirq.measure(q1)
+        cirq.XPowGate(dimension=3)(q0) ** 2,
+        TestChannel()(q0),
+        TestChannel()(q0),
+        cirq.measure(q0),
+        cirq.measure(q1),
     )
 
     simulator = cirq.DensityMatrixSimulator(dtype=dtype, split_untangled_states=split)
@@ -306,7 +298,10 @@ def test_run_qudits_repetitions_measure_at_end(dtype: Type[np.number], split: bo
         for b0 in [0, 1]:
             for b1 in [0, 1, 2]:
                 circuit = cirq.Circuit(
-                    (cirq.X**b0)(q0), PlusGate(3, b1)(q1), cirq.measure(q0), cirq.measure(q1)
+                    (cirq.X**b0)(q0),
+                    cirq.XPowGate(dimension=3)(q1) ** b1,
+                    cirq.measure(q0),
+                    cirq.measure(q1),
                 )
                 result = simulator.run(circuit, repetitions=3)
                 np.testing.assert_equal(
@@ -372,11 +367,11 @@ def test_run_qudits_repetitions_measurement_not_terminal(dtype: Type[np.number],
             for b1 in [0, 1, 2]:
                 circuit = cirq.Circuit(
                     (cirq.X**b0)(q0),
-                    PlusGate(3, b1)(q1),
+                    cirq.XPowGate(dimension=3)(q1) ** b1,
                     cirq.measure(q0),
                     cirq.measure(q1),
                     cirq.H(q0),
-                    PlusGate(3, -b1)(q1),
+                    cirq.XPowGate(dimension=3)(q1) ** (-b1),
                 )
                 result = simulator.run(circuit, repetitions=3)
                 np.testing.assert_equal(
@@ -437,7 +432,9 @@ def test_run_measure_multiple_qudits(dtype: Type[np.number], split: bool):
     simulator = cirq.DensityMatrixSimulator(dtype=dtype, split_untangled_states=split)
     for b0 in [0, 1]:
         for b1 in [0, 1, 2]:
-            circuit = cirq.Circuit((cirq.X**b0)(q0), PlusGate(3, b1)(q1), cirq.measure(q0, q1))
+            circuit = cirq.Circuit(
+                (cirq.X**b0)(q0), cirq.XPowGate(dimension=3)(q1) ** b1, cirq.measure(q0, q1)
+            )
             result = simulator.run(circuit, repetitions=3)
             np.testing.assert_equal(result.measurements, {'0 (d=2),1 (d=3)': [[b0, b1]] * 3})
 
@@ -497,7 +494,7 @@ def test_simulate(dtype: Type[np.number], split: bool):
 def test_simulate_qudits(dtype: Type[np.number], split: bool):
     q0, q1 = cirq.LineQid.for_qid_shape((2, 3))
     simulator = cirq.DensityMatrixSimulator(dtype=dtype, split_untangled_states=split)
-    circuit = cirq.Circuit(cirq.H(q0), PlusGate(3, 2)(q1))
+    circuit = cirq.Circuit(cirq.H(q0), cirq.XPowGate(dimension=3)(q1) ** 2)
     result = simulator.simulate(circuit, qubit_order=[q1, q0])
     expected = np.zeros((6, 6))
     expected[4:, 4:] = np.ones((2, 2)) / 2
@@ -566,13 +563,16 @@ def test_simulate_qudit_increments(dtype: Type[np.number], split: bool):
     for b0 in [0, 1]:
         for b1 in [0, 1, 2]:
             circuit = cirq.Circuit(
-                (cirq.X**b0)(q0), (PlusGate(3)(q1),) * b1, cirq.measure(q0), cirq.measure(q1)
+                (cirq.X**b0)(q0),
+                (cirq.XPowGate(dimension=3)(q1),) * b1,
+                cirq.measure(q0),
+                cirq.measure(q1),
             )
             result = simulator.simulate(circuit)
             np.testing.assert_equal(result.measurements, {'0 (d=2)': [b0], '1 (d=3)': [b1]})
             expected_density_matrix = np.zeros(shape=(6, 6))
             expected_density_matrix[b0 * 3 + b1, b0 * 3 + b1] = 1.0
-            np.testing.assert_equal(result.final_density_matrix, expected_density_matrix)
+            np.testing.assert_allclose(result.final_density_matrix, expected_density_matrix)
 
 
 @pytest.mark.parametrize('dtype', [np.complex64, np.complex128])
@@ -623,13 +623,17 @@ def test_simulate_initial_qudit_state(dtype: Type[np.number], split: bool):
     simulator = cirq.DensityMatrixSimulator(dtype=dtype, split_untangled_states=split)
     for b0 in [0, 1, 2]:
         for b1 in [0, 1, 2, 3]:
-            circuit = cirq.Circuit(PlusGate(3, b0)(q0), PlusGate(4, b1)(q1))
+            circuit = cirq.Circuit(
+                cirq.XPowGate(dimension=3)(q0) ** b0, cirq.XPowGate(dimension=4)(q1) ** b1
+            )
             result = simulator.simulate(circuit, initial_state=6)
             expected_density_matrix = np.zeros(shape=(12, 12))
             expected_density_matrix[
                 (b0 + 1) % 3 * 4 + (b1 + 2) % 4, (b0 + 1) % 3 * 4 + (b1 + 2) % 4
             ] = 1.0
-            np.testing.assert_equal(result.final_density_matrix, expected_density_matrix)
+            np.testing.assert_allclose(
+                result.final_density_matrix, expected_density_matrix, atol=1e-15
+            )
 
 
 @pytest.mark.parametrize('dtype', [np.complex64, np.complex128])
@@ -684,7 +688,9 @@ def test_simulate_measure_multiple_qudits(dtype: Type[np.number], split: bool):
     simulator = cirq.DensityMatrixSimulator(dtype=dtype, split_untangled_states=split)
     for b0 in [0, 1]:
         for b1 in [0, 1, 2]:
-            circuit = cirq.Circuit((cirq.X**b0)(q0), PlusGate(3, b1)(q1), cirq.measure(q0, q1))
+            circuit = cirq.Circuit(
+                (cirq.X**b0)(q0), cirq.XPowGate(dimension=3)(q1) ** b1, cirq.measure(q0, q1)
+            )
             result = simulator.simulate(circuit)
             np.testing.assert_equal(result.measurements, {'0 (d=2),1 (d=3)': [b0, b1]})
 
@@ -735,7 +741,10 @@ def test_simulate_moment_steps(dtype: Type[np.number], split: bool):
 def test_simulate_moment_steps_qudits(dtype: Type[np.number], split: bool):
     q0, q1 = cirq.LineQid.for_qid_shape((2, 3))
     circuit = cirq.Circuit(
-        PlusGate(2, 1)(q0), PlusGate(3, 1)(q1), cirq.reset(q1), PlusGate(3, 1)(q1)
+        cirq.XPowGate(dimension=2)(q0),
+        cirq.XPowGate(dimension=3)(q1),
+        cirq.reset(q1),
+        cirq.XPowGate(dimension=3)(q1),
     )
     simulator = cirq.DensityMatrixSimulator(dtype=dtype, split_untangled_states=split)
     for i, step in enumerate(simulator.simulate_moment_steps(circuit)):
@@ -1400,29 +1409,12 @@ def test_random_seed_non_terminal_measurements_deterministic():
 
 
 def test_simulate_with_invert_mask():
-    class PlusGate(cirq.Gate):
-        """A qudit gate that increments a qudit state mod its dimension."""
-
-        def __init__(self, dimension, increment=1):
-            self.dimension = dimension
-            self.increment = increment % dimension
-
-        def _qid_shape_(self):
-            return (self.dimension,)
-
-        def _unitary_(self):
-            inc = (self.increment - 1) % self.dimension + 1
-            u = np.empty((self.dimension, self.dimension))
-            u[inc:] = np.eye(self.dimension)[:-inc]
-            u[:inc] = np.eye(self.dimension)[-inc:]
-            return u
-
     q0, q1, q2, q3, q4 = cirq.LineQid.for_qid_shape((2, 3, 3, 3, 4))
     c = cirq.Circuit(
-        PlusGate(2, 1)(q0),
-        PlusGate(3, 1)(q2),
-        PlusGate(3, 2)(q3),
-        PlusGate(4, 3)(q4),
+        cirq.XPowGate(dimension=2)(q0),
+        cirq.XPowGate(dimension=3)(q2),
+        cirq.XPowGate(dimension=3)(q3) ** 2,
+        cirq.XPowGate(dimension=4)(q4) ** 3,
         cirq.measure(q0, q1, q2, q3, q4, key='a', invert_mask=(True,) * 4),
     )
     assert np.all(cirq.DensityMatrixSimulator().run(c).measurements['a'] == [[0, 1, 0, 2, 3]])
diff --git a/cirq-core/cirq/sim/simulator_test.py b/cirq-core/cirq/sim/simulator_test.py
--- a/cirq-core/cirq/sim/simulator_test.py
+++ b/cirq-core/cirq/sim/simulator_test.py
@@ -389,29 +389,12 @@ def test_sample_repeated_measurement_keys():
 
 
 def test_simulate_with_invert_mask():
-    class PlusGate(cirq.Gate):
-        """A qudit gate that increments a qudit state mod its dimension."""
-
-        def __init__(self, dimension, increment=1):
-            self.dimension = dimension
-            self.increment = increment % dimension
-
-        def _qid_shape_(self):
-            return (self.dimension,)
-
-        def _unitary_(self):
-            inc = (self.increment - 1) % self.dimension + 1
-            u = np.empty((self.dimension, self.dimension))
-            u[inc:] = np.eye(self.dimension)[:-inc]
-            u[:inc] = np.eye(self.dimension)[-inc:]
-            return u
-
     q0, q1, q2, q3, q4 = cirq.LineQid.for_qid_shape((2, 3, 3, 3, 4))
     c = cirq.Circuit(
-        PlusGate(2, 1)(q0),
-        PlusGate(3, 1)(q2),
-        PlusGate(3, 2)(q3),
-        PlusGate(4, 3)(q4),
+        cirq.XPowGate(dimension=2)(q0),
+        cirq.XPowGate(dimension=3)(q2),
+        cirq.XPowGate(dimension=3)(q3) ** 2,
+        cirq.XPowGate(dimension=4)(q4) ** 3,
         cirq.measure(q0, q1, q2, q3, q4, key='a', invert_mask=(True,) * 4),
     )
     assert np.all(cirq.Simulator().run(c).measurements['a'] == [[0, 1, 0, 2, 3]])
diff --git a/cirq-core/cirq/sim/sparse_simulator_test.py b/cirq-core/cirq/sim/sparse_simulator_test.py
--- a/cirq-core/cirq/sim/sparse_simulator_test.py
+++ b/cirq-core/cirq/sim/sparse_simulator_test.py
@@ -64,7 +64,7 @@ def test_run_reset(dtype: Type[np.number], split: bool):
     simulator = cirq.Simulator(dtype=dtype, split_untangled_states=split)
     circuit = cirq.Circuit(
         cirq.H(q0),
-        PlusGate(3, 2)(q1),
+        cirq.XPowGate(dimension=3)(q1) ** 2,
         cirq.reset(q0),
         cirq.measure(q0, key='m0'),
         cirq.measure(q1, key='m1a'),
@@ -352,24 +352,6 @@ def test_simulate(dtype: Type[np.number], split: bool):
     assert len(result.measurements) == 0
 
 
-class PlusGate(cirq.Gate):
-    """A qudit gate that increments a qudit state mod its dimension."""
-
-    def __init__(self, dimension, increment=1):
-        self.dimension = dimension
-        self.increment = increment % dimension
-
-    def _qid_shape_(self):
-        return (self.dimension,)
-
-    def _unitary_(self):
-        inc = (self.increment - 1) % self.dimension + 1
-        u = np.empty((self.dimension, self.dimension))
-        u[inc:] = np.eye(self.dimension)[:-inc]
-        u[:inc] = np.eye(self.dimension)[-inc:]
-        return u
-
-
 class _TestMixture(cirq.Gate):
     def __init__(self, gate_options):
         self.gate_options = gate_options
@@ -386,7 +368,7 @@ def _mixture_(self):
 def test_simulate_qudits(dtype: Type[np.number], split: bool):
     q0, q1 = cirq.LineQid.for_qid_shape((3, 4))
     simulator = cirq.Simulator(dtype=dtype, split_untangled_states=split)
-    circuit = cirq.Circuit(PlusGate(3)(q0), PlusGate(4, increment=3)(q1))
+    circuit = cirq.Circuit(cirq.XPowGate(dimension=3)(q0), cirq.XPowGate(dimension=4)(q1) ** 3)
     result = simulator.simulate(circuit, qubit_order=[q0, q1])
     expected = np.zeros(12)
     expected[4 * 1 + 3] = 1
@@ -417,7 +399,13 @@ def test_simulate_mixtures(dtype: Type[np.number], split: bool):
 def test_simulate_qudit_mixtures(dtype: Type[np.number], split: bool):
     q0 = cirq.LineQid(0, 3)
     simulator = cirq.Simulator(dtype=dtype, split_untangled_states=split)
-    mixture = _TestMixture([PlusGate(3, 0), PlusGate(3, 1), PlusGate(3, 2)])
+    mixture = _TestMixture(
+        [
+            cirq.XPowGate(dimension=3) ** 0,
+            cirq.XPowGate(dimension=3),
+            cirq.XPowGate(dimension=3) ** 2,
+        ]
+    )
     circuit = cirq.Circuit(mixture(q0), cirq.measure(q0))
     counts = {0: 0, 1: 0, 2: 0}
     for _ in range(300):
@@ -834,10 +822,10 @@ def test_measure_at_end_invert_mask_partial():
 def test_qudit_invert_mask():
     q0, q1, q2, q3, q4 = cirq.LineQid.for_qid_shape((2, 3, 3, 3, 4))
     c = cirq.Circuit(
-        PlusGate(2, 1)(q0),
-        PlusGate(3, 1)(q2),
-        PlusGate(3, 2)(q3),
-        PlusGate(4, 3)(q4),
+        cirq.XPowGate(dimension=2)(q0),
+        cirq.XPowGate(dimension=3)(q2),
+        cirq.XPowGate(dimension=3)(q3) ** 2,
+        cirq.XPowGate(dimension=4)(q4) ** 3,
         cirq.measure(q0, q1, q2, q3, q4, key='a', invert_mask=(True,) * 4),
     )
     assert np.all(cirq.Simulator().run(c).measurements['a'] == [[0, 1, 0, 2, 3]])