Repository: quantumlib/Cirq

Issue:
ArithmeticOperation apply unitary assumes reshape always returns a view
This is not true. In some cases reshape returns a copy. This can result in undefined data going into the result, because the workspace buffer is never initialized.

The safe way to reshape and get an error if it's not a mutable view is to make an identical view `view = buffer[...]` and then *set* its shape `view.shape = shape`.

The code needs to be fixed to work in the case where reshape would fail, and a test needs to be added to trigger this case. I ran into this as part of https://github.com/quantumlib/Cirq/pull/2137 .

Repro:

```
def test_reshape_referencing():
    class Op1(cirq.ArithmeticOperation):
        def apply(self, *register_values: int):
            return register_values[0] + 1

        def registers(self):
            return [cirq.LineQubit.range(2)[::-1]]

        def with_registers(self, *new_registers):
            raise NotImplementedError()

    state = np.ones(4, dtype=np.complex64) / 2
    output = cirq.final_wavefunction(
        cirq.Circuit.from_ops(Op1()),
        initial_state=state)
    np.testing.assert_allclose(state, output)
```

Results in:

```
E       AssertionError: 
E       Not equal to tolerance rtol=1e-07, atol=0
E       
E       Mismatch: 100%
E       Max absolute difference: 0.5
E       Max relative difference: inf
E        x: array([0.5+0.j, 0.5+0.j, 0.5+0.j, 0.5+0.j], dtype=complex64)
E        y: array([0.+0.j, 1.+0.j, 0.+0.j, 0.+0.j], dtype=complex64)
```

But any permutation of the all-plus initial state should be the all-plus state.


Discussion:

Resolving patch:
diff --git a/cirq/ops/arithmetic_operation.py b/cirq/ops/arithmetic_operation.py
--- a/cirq/ops/arithmetic_operation.py
+++ b/cirq/ops/arithmetic_operation.py
@@ -223,7 +223,11 @@ def _apply_unitary_(self, args: 'cirq.ApplyUnitaryArgs'):
             cast(List[Union[int, slice]], inputs).append(slice(None))
             dst[outputs] = src[inputs]
 
-        return args.available_buffer
+        # In case the reshaped arrays were copies instead of views.
+        dst.shape = transposed_args.available_buffer.shape
+        transposed_args.target_tensor[...] = dst
+
+        return args.target_tensor
 
 
 def _describe_bad_arithmetic_changed_const(
diff --git a/cirq/schedules/schedule.py b/cirq/schedules/schedule.py
--- a/cirq/schedules/schedule.py
+++ b/cirq/schedules/schedule.py
@@ -203,6 +203,9 @@ def to_circuit(self) -> Circuit:
                                strategy=InsertStrategy.INLINE)
         return circuit
 
+    def _qid_shape_(self):
+        return protocols.qid_shape(self.to_circuit())
+
     def _has_unitary_(self):
         return protocols.has_unitary(self.to_circuit())
 


Test patch:
diff --git a/cirq/ops/arithmetic_operation_test.py b/cirq/ops/arithmetic_operation_test.py
--- a/cirq/ops/arithmetic_operation_test.py
+++ b/cirq/ops/arithmetic_operation_test.py
@@ -114,6 +114,9 @@ def apply(self, target_value, input_value):
 
     np.testing.assert_allclose(cirq.unitary(Add(1, 0)), np.eye(1))
 
+    cirq.testing.assert_has_consistent_apply_unitary(
+        Add(cirq.LineQubit.range(2), cirq.LineQubit.range(2)))
+
 
 def test_arithmetic_operation_qubits():
 
@@ -149,3 +152,22 @@ def apply(self, target_value, input_value):
     op4 = op3.with_qubits(q0, q1, q2, q3)
     assert op4.registers() == ([q0, q1, q2], [q3], 1)
     assert op4.qubits == (q0, q1, q2, q3)
+
+
+def test_reshape_referencing():
+
+    class Op1(cirq.ArithmeticOperation):
+
+        def apply(self, *register_values: int):
+            return register_values[0] + 1
+
+        def registers(self):
+            return [cirq.LineQubit.range(2)[::-1]]
+
+        def with_registers(self, *new_registers):
+            raise NotImplementedError()
+
+    state = np.ones(4, dtype=np.complex64) / 2
+    output = cirq.final_wavefunction(cirq.Circuit.from_ops(Op1()),
+                                     initial_state=state)
+    np.testing.assert_allclose(state, output)
diff --git a/cirq/schedules/schedule_test.py b/cirq/schedules/schedule_test.py
--- a/cirq/schedules/schedule_test.py
+++ b/cirq/schedules/schedule_test.py
@@ -279,7 +279,7 @@ def test_unitary():
     schedule = cirq.Schedule(device=UNCONSTRAINED_DEVICE,
                              scheduled_operations=[op])
 
-    cirq.testing.assert_has_consistent_apply_unitary(schedule, qubit_count=1)
+    cirq.testing.assert_has_consistent_apply_unitary(schedule)
     np.testing.assert_allclose(cirq.unitary(schedule), cirq.unitary(cirq.H))
     assert cirq.has_unitary(schedule)
 
diff --git a/cirq/testing/__init__.py b/cirq/testing/__init__.py
--- a/cirq/testing/__init__.py
+++ b/cirq/testing/__init__.py
@@ -22,7 +22,6 @@
 
 from cirq.testing.circuit_compare import (
     assert_circuits_with_terminal_measurements_are_equivalent,
-    assert_eigen_gate_has_consistent_apply_unitary,
     assert_has_consistent_apply_unitary,
     assert_has_consistent_apply_unitary_for_various_exponents,
     assert_has_diagram,
diff --git a/cirq/testing/circuit_compare.py b/cirq/testing/circuit_compare.py
--- a/cirq/testing/circuit_compare.py
+++ b/cirq/testing/circuit_compare.py
@@ -11,15 +11,17 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-from typing import Any, Dict, Iterable, List, Optional, Sequence, Type, cast
+from typing import Any, Dict, Iterable, List, Optional, Sequence
 
 from collections import defaultdict
 import itertools
+import random
+
 import numpy as np
 import sympy
 
 from cirq import circuits, ops, linalg, protocols
+from cirq.testing import lin_alg_utils
 
 
 def highlight_text_differences(actual: str, expected: str) -> str:
@@ -241,11 +243,8 @@ def assert_has_diagram(
     )
 
 
-def assert_has_consistent_apply_unitary(
-        val: Any,
-        *,
-        qubit_count: Optional[int] = None,
-        atol: float=1e-8) -> None:
+def assert_has_consistent_apply_unitary(val: Any, *,
+                                        atol: float = 1e-8) -> None:
     """Tests whether a value's _apply_unitary_ is correct.
 
     Contrasts the effects of the value's `_apply_unitary_` with the
@@ -253,42 +252,23 @@ def assert_has_consistent_apply_unitary(
 
     Args:
         val: The value under test. Should have a `__pow__` method.
-        qubit_count: Usually inferred. The number of qubits the value acts on.
-            This argument isn't needed if the gate has a unitary matrix or
-            implements `cirq.SingleQubitGate`/`cirq.TwoQubitGate`/
-            `cirq.ThreeQubitGate`.
         atol: Absolute error tolerance.
     """
 
+    _assert_apply_unitary_works_when_axes_transposed(val, atol=atol)
+
     expected = protocols.unitary(val, default=None)
 
-    qubit_counts = [
-        qubit_count,
-        # Only fall back to using the unitary size if num_qubits or qid_shape
-        # protocols are not defined.
-        protocols.num_qubits(val,
-                             default=expected.shape[0].bit_length() -
-                             1 if expected is not None else None),
-        _infer_qubit_count(val)
-    ]
-    qubit_counts = [e for e in qubit_counts if e is not None]
-    if not qubit_counts:
-        raise NotImplementedError(
-            'Failed to infer qubit count of <{!r}>. Specify it.'.format(
-                val))
-    assert len(set(qubit_counts)) == 1, (
-        'Inconsistent qubit counts from different methods: {}'.format(
-            qubit_counts))
-    n = cast(int, qubit_counts[0])
-    qid_shape = protocols.qid_shape(val, default=(2,) * n)
+    qid_shape = protocols.qid_shape(val)
 
     eye = linalg.eye_tensor((2,) + qid_shape, dtype=np.complex128)
     actual = protocols.apply_unitary(
         unitary_value=val,
-        args=protocols.ApplyUnitaryArgs(
-            target_tensor=eye,
-            available_buffer=np.ones_like(eye) * float('nan'),
-            axes=list(range(1, n + 1))),
+        args=protocols.ApplyUnitaryArgs(target_tensor=eye,
+                                        available_buffer=np.ones_like(eye) *
+                                        float('nan'),
+                                        axes=list(range(1,
+                                                        len(qid_shape) + 1))),
         default=None)
 
     # If you don't have a unitary, you shouldn't be able to apply a unitary.
@@ -305,42 +285,76 @@ def assert_has_consistent_apply_unitary(
                                    atol=atol)
 
 
-def assert_eigen_gate_has_consistent_apply_unitary(
-        eigen_gate_type: Type[ops.EigenGate],
-        *,
-        exponents=(0, 1, -1, 0.5, 0.25, -0.5, 0.1, sympy.Symbol('s')),
-        global_shifts=(0, 0.5, -0.5, 0.1),
-        qubit_count: Optional[int] = None) -> None:
-    """Tests whether an EigenGate type's _apply_unitary_ is correct.
+def _assert_apply_unitary_works_when_axes_transposed(val: Any,
+                                                     *,
+                                                     atol: float = 1e-8
+                                                    ) -> None:
+    """Tests whether a value's _apply_unitary_ handles out-of-order axes.
 
-    Contrasts the effects of the gate's `_apply_unitary_` with the
-    matrix returned by the gate's `_unitary_` method, trying various values for
-    the gate exponent and global shift.
+    A common mistake to make when implementing `_apply_unitary_` is to assume
+    that the incoming axes will be contiguous, or ascending, or that they can be
+    flattened, or that other axes have a length of two, etc, etc ,etc. This
+    method checks that `_apply_unitary_` does the same thing to out-of-order
+    axes that it does to contiguous in-order axes.
 
     Args:
-        eigen_gate_type: The type of gate to test. The type must have an
-            __init__ method that takes an exponent and a global_shift.
-        exponents: The exponents to try. Defaults to a variety of special and
-            arbitrary angles, as well as a parameterized angle (a symbol).
-        global_shifts: The global shifts to try. Defaults to a variety of
-            special angles.
-        qubit_count: The qubit count to use for the gate. This argument isn't
-            needed if the gate has a unitary matrix or implements
-            `cirq.SingleQubitGate`/`cirq.TwoQubitGate`/`cirq.ThreeQubitGate`; it
-            will be inferred.
+        val: The operation, gate, or other unitary object to test.
+        atol: Absolute error tolerance.
     """
-    for exponent in exponents:
-        for shift in global_shifts:
-            assert_has_consistent_apply_unitary(
-                eigen_gate_type(exponent=exponent, global_shift=shift),
-                qubit_count=qubit_count)
+
+    # Only test custom apply unitary methods.
+    if not hasattr(val, '_apply_unitary_') or not protocols.has_unitary(val):
+        return
+
+    # Pick sizes and shapes.
+    shape = protocols.qid_shape(val)
+    n = len(shape)
+    padded_shape = shape + (1, 2, 2, 3)
+    padded_n = len(padded_shape)
+    size = np.product(padded_shape).item()
+
+    # Shuffle the axes.
+    permutation = list(range(padded_n))
+    random.shuffle(permutation)
+    transposed_shape = [0] * padded_n
+    for i in range(padded_n):
+        transposed_shape[permutation[i]] = padded_shape[i]
+
+    # Prepare input states.
+    in_order_input = lin_alg_utils.random_superposition(size).reshape(
+        padded_shape)
+    out_of_order_input = np.empty(shape=transposed_shape, dtype=np.complex128)
+    out_of_order_input.transpose(permutation)[...] = in_order_input
+
+    # Apply to in-order and out-of-order axes.
+    in_order_output = protocols.apply_unitary(
+        val,
+        protocols.ApplyUnitaryArgs(
+            target_tensor=in_order_input,
+            available_buffer=np.empty_like(in_order_input),
+            axes=range(n)))
+    out_of_order_output = protocols.apply_unitary(
+        val,
+        protocols.ApplyUnitaryArgs(
+            target_tensor=out_of_order_input,
+            available_buffer=np.empty_like(out_of_order_input),
+            axes=permutation[:n]))
+
+    # Put the out of order output back into order, to enable comparison.
+    reordered_output = out_of_order_output.transpose(permutation)
+
+    # The results should be identical.
+    if not np.allclose(in_order_output, reordered_output, atol=atol):
+        raise AssertionError(
+            f'The _apply_unitary_ method of {repr(val)} acted differently on '
+            f'out-of-order axes than on in-order axes.\n'
+            f'\n'
+            f'The failing axis order: {repr(permutation[:n])}')
 
 
 def assert_has_consistent_apply_unitary_for_various_exponents(
-        val: Any,
-        *,
-        exponents=(0, 1, -1, 0.5, 0.25, -0.5, 0.1, sympy.Symbol('s')),
-        qubit_count: Optional[int] = None) -> None:
+        val: Any, *,
+        exponents=(0, 1, -1, 0.5, 0.25, -0.5, 0.1, sympy.Symbol('s'))) -> None:
     """Tests whether a value's _apply_unitary_ is correct.
 
     Contrasts the effects of the value's `_apply_unitary_` with the
@@ -353,21 +367,14 @@ def assert_has_consistent_apply_unitary_for_various_exponents(
             arbitrary angles, as well as a parameterized angle (a symbol). If
             the value's `__pow__` returns `NotImplemented` for any of these,
             they are skipped.
-        qubit_count: A minimum qubit count for the test system. This argument
-            isn't needed if the gate has a unitary matrix or implements
-            `cirq.SingleQubitGate`/`cirq.TwoQubitGate`/`cirq.ThreeQubitGate`; it
-            will be inferred.
     """
     for exponent in exponents:
         gate = protocols.pow(val, exponent, default=None)
         if gate is not None:
-            assert_has_consistent_apply_unitary(
-                gate,
-                qubit_count=qubit_count)
+            assert_has_consistent_apply_unitary(gate)
 
 
-def assert_has_consistent_qid_shape(val: Any,
-                                    qubit_count: Optional[int] = None) -> None:
+def assert_has_consistent_qid_shape(val: Any) -> None:
     """Tests whether a value's `_qid_shape_` and `_num_qubits_` are correct and
     consistent.
 
@@ -378,7 +385,6 @@ def assert_has_consistent_qid_shape(val: Any,
     Args:
         val: The value under test. Should have `_qid_shape_` and/or
             `num_qubits_` methods. Can optionally have a `qubits` property.
-        qubit_count: The expected number of qubits val should use.
     """
     default = (-1,)
     qid_shape = protocols.qid_shape(val, default)
@@ -390,20 +396,8 @@ def assert_has_consistent_qid_shape(val: Any,
     assert len(qid_shape) == num_qubits, (
         f'Length of qid_shape and num_qubits disagree: {qid_shape}, '
         f'{num_qubits}')
-    if qubit_count is not None:
-        assert qubit_count == num_qubits, (
-            f'Expected qubits and num_qubits disagree: {qubit_count}, '
-            f'{num_qubits}')
-    infer_qubit_count = _infer_qubit_count(val)
-    if infer_qubit_count is not None:
-        assert infer_qubit_count == num_qubits, (
-            f'Length of qubits and num_qubits disagree: {infer_qubit_count}, '
-            f'{num_qubits}')
-
-
-def _infer_qubit_count(val: Any) -> Optional[int]:
+
     if isinstance(val, ops.Operation):
-        return len(val.qubits)
-    if isinstance(val, ops.Gate):
-        return protocols.num_qubits(val)
-    return None
+        assert num_qubits == len(val.qubits), (
+            f'Length of num_qubits and val.qubits disagrees: {num_qubits}, '
+            f'{len(val.qubits)}')
diff --git a/cirq/testing/circuit_compare_test.py b/cirq/testing/circuit_compare_test.py
--- a/cirq/testing/circuit_compare_test.py
+++ b/cirq/testing/circuit_compare_test.py
@@ -17,6 +17,8 @@
 import numpy as np
 
 import cirq
+from cirq.testing.circuit_compare import (
+    _assert_apply_unitary_works_when_axes_transposed,)
 
 
 def test_sensitive_to_phase():
@@ -313,33 +315,45 @@ def _apply_unitary_(self, args: cirq.ApplyUnitaryArgs) -> np.ndarray:
         def _unitary_(self):
             return np.eye(2)
 
+        def _num_qubits_(self):
+            return 1
+
     with pytest.raises(AssertionError):
         cirq.testing.assert_has_consistent_apply_unitary(
             IdentityReturningUnalteredWorkspace())
 
     class DifferentEffect:
         def _apply_unitary_(self, args: cirq.ApplyUnitaryArgs) -> np.ndarray:
-            args.available_buffer[0] = args.target_tensor[1]
-            args.available_buffer[1] = args.target_tensor[0]
+            o = args.subspace_index(0)
+            i = args.subspace_index(1)
+            args.available_buffer[o] = args.target_tensor[i]
+            args.available_buffer[i] = args.target_tensor[o]
             return args.available_buffer
 
         def _unitary_(self):
             return np.eye(2, dtype=np.complex128)
 
+        def _num_qubits_(self):
+            return 1
+
     with pytest.raises(AssertionError):
         cirq.testing.assert_has_consistent_apply_unitary(
             DifferentEffect())
 
     class IgnoreAxisEffect:
         def _apply_unitary_(self, args: cirq.ApplyUnitaryArgs) -> np.ndarray:
-            args.available_buffer[0] = args.target_tensor[1]
-            args.available_buffer[1] = args.target_tensor[0]
+            if args.target_tensor.shape[0] > 1:
+                args.available_buffer[0] = args.target_tensor[1]
+                args.available_buffer[1] = args.target_tensor[0]
             return args.available_buffer
 
         def _unitary_(self):
             return np.array([[0, 1], [1, 0]])
 
-    with pytest.raises(AssertionError, match='Not equal'):
+        def _num_qubits_(self):
+            return 1
+
+    with pytest.raises(AssertionError, match='Not equal|acted differently'):
         cirq.testing.assert_has_consistent_apply_unitary(
             IgnoreAxisEffect())
 
@@ -354,6 +368,9 @@ def _apply_unitary_(self, args: cirq.ApplyUnitaryArgs) -> np.ndarray:
         def _unitary_(self):
             return np.array([[0, 1], [1, 0]])
 
+        def _num_qubits_(self):
+            return 1
+
     cirq.testing.assert_has_consistent_apply_unitary(
         SameEffect())
 
@@ -380,6 +397,9 @@ def __init__(self, power):
         def __pow__(self, power):
             return BadExponent(self.power * power)
 
+        def _num_qubits_(self):
+            return 1
+
         def _apply_unitary_(self, args: cirq.ApplyUnitaryArgs) -> np.ndarray:
             i = args.subspace_index(1)
             args.target_tensor[i] *= self.power * 2
@@ -393,35 +413,32 @@ def _unitary_(self):
 
     with pytest.raises(AssertionError):
         cirq.testing.assert_has_consistent_apply_unitary_for_various_exponents(
-            BadExponent(1),
-            exponents=[1, 2],
-            qubit_count=1)
+            BadExponent(1), exponents=[1, 2])
 
     class EffectWithoutUnitary:
+
+        def _num_qubits_(self):
+            return 1
+
         def _apply_unitary_(self, args: cirq.ApplyUnitaryArgs) -> np.ndarray:
             return args.target_tensor
 
-    with pytest.raises(AssertionError):
-        cirq.testing.assert_has_consistent_apply_unitary(
-            EffectWithoutUnitary(),
-            qubit_count=1)
+    cirq.testing.assert_has_consistent_apply_unitary(EffectWithoutUnitary())
 
     class NoEffect:
+
+        def _num_qubits_(self):
+            return 1
+
         def _apply_unitary_(self, args: cirq.ApplyUnitaryArgs) -> np.ndarray:
             return NotImplemented
 
-    cirq.testing.assert_has_consistent_apply_unitary(
-        NoEffect(),
-        qubit_count=1)
+    cirq.testing.assert_has_consistent_apply_unitary(NoEffect())
 
     class UnknownCountEffect:
         pass
 
-    cirq.testing.assert_has_consistent_apply_unitary(
-        UnknownCountEffect(),
-        qubit_count=1)
-
-    with pytest.raises(NotImplementedError):
+    with pytest.raises(TypeError, match="no _num_qubits_ or _qid_shape_"):
         cirq.testing.assert_has_consistent_apply_unitary(
             UnknownCountEffect())
 
@@ -432,12 +449,6 @@ class UnknownCountEffect:
         cirq.X.on(cirq.NamedQubit('q')))
 
 
-def test_inconsistent_qubit_count():
-    with pytest.raises(AssertionError, match='Inconsistent'):
-        cirq.testing.assert_has_consistent_apply_unitary(
-            cirq.X, qubit_count=2)
-
-
 def test_assert_has_consistent_qid_shape():
 
     class ConsistentGate(cirq.Gate):
@@ -446,7 +457,7 @@ def _num_qubits_(self):
             return 4
 
         def _qid_shape_(self):
-            return (1, 2, 3, 4)
+            return 1, 2, 3, 4
 
     class InconsistentGate(cirq.Gate):
 
@@ -454,7 +465,7 @@ def _num_qubits_(self):
             return 2
 
         def _qid_shape_(self):
-            return (1, 2, 3, 4)
+            return 1, 2, 3, 4
 
     class BadShapeGate(cirq.Gate):
 
@@ -462,7 +473,7 @@ def _num_qubits_(self):
             return 4
 
         def _qid_shape_(self):
-            return (1, 2, 0, 4)
+            return 1, 2, 0, 4
 
     class ConsistentOp(cirq.Operation):
 
@@ -525,7 +536,7 @@ def _num_qubits_(self):
             return 4  # coverage: ignore
 
         def _qid_shape_(self):
-            return (1, 2)
+            return 1, 2
 
     class NoProtocol:
         pass
@@ -543,3 +554,38 @@ class NoProtocol:
     with pytest.raises(AssertionError, match='disagree'):
         cirq.testing.assert_has_consistent_qid_shape(InconsistentOp3())
     cirq.testing.assert_has_consistent_qid_shape(NoProtocol())
+
+
+def test_assert_apply_unitary_works_when_axes_transposed_failure():
+
+    class BadOp:
+
+        def _apply_unitary_(self, args: cirq.ApplyUnitaryArgs):
+            # Get a more convenient view of the data.
+            a, b = args.axes
+            rest = list(range(len(args.target_tensor.shape)))
+            rest.remove(a)
+            rest.remove(b)
+            size = args.target_tensor.size
+            view = args.target_tensor.transpose([a, b, *rest])
+            view = view.reshape((4, size // 4))  # Oops. Reshape might copy.
+
+            # Apply phase gradient.
+            view[1, ...] *= 1j
+            view[2, ...] *= -1
+            view[3, ...] *= -1j
+            return args.target_tensor
+
+        def _num_qubits_(self):
+            return 2
+
+    bad_op = BadOp()
+    assert cirq.has_unitary(bad_op)
+
+    # Appears to work.
+    np.testing.assert_allclose(cirq.unitary(bad_op), np.diag([1, 1j, -1, -1j]))
+    # But fails the more discerning test.
+    with pytest.raises(AssertionError,
+                       match='acted differently on out-of-order axes'):
+        for _ in range(100):  # Axis orders chosen at random. Brute force a hit.
+            _assert_apply_unitary_works_when_axes_transposed(bad_op)
diff --git a/cirq/testing/consistent_protocols.py b/cirq/testing/consistent_protocols.py
--- a/cirq/testing/consistent_protocols.py
+++ b/cirq/testing/consistent_protocols.py
@@ -47,21 +47,20 @@ def assert_implements_consistent_protocols(
     local_vals = local_vals or {}
 
     _assert_meets_standards_helper(val,
-                                   qubit_count,
-                                   ignoring_global_phase,
-                                   setup_code,
-                                   global_vals,
-                                   local_vals)
+                                   ignoring_global_phase=ignoring_global_phase,
+                                   setup_code=setup_code,
+                                   global_vals=global_vals,
+                                   local_vals=local_vals)
 
     for exponent in exponents:
         p = protocols.pow(val, exponent, None)
         if p is not None:
-            _assert_meets_standards_helper(val**exponent,
-                                           qubit_count,
-                                           ignoring_global_phase,
-                                           setup_code,
-                                           global_vals,
-                                           local_vals)
+            _assert_meets_standards_helper(
+                val**exponent,
+                ignoring_global_phase=ignoring_global_phase,
+                setup_code=setup_code,
+                global_vals=global_vals,
+                local_vals=local_vals)
 
 
 def assert_eigengate_implements_consistent_protocols(
@@ -80,12 +79,11 @@ def assert_eigengate_implements_consistent_protocols(
     for exponent in exponents:
         for shift in global_shifts:
             _assert_meets_standards_helper(
-                    eigen_gate_type(exponent=exponent, global_shift=shift),
-                    qubit_count,
-                    ignoring_global_phase,
-                    setup_code,
-                    global_vals,
-                    local_vals)
+                eigen_gate_type(exponent=exponent, global_shift=shift),
+                ignoring_global_phase=ignoring_global_phase,
+                setup_code=setup_code,
+                global_vals=global_vals,
+                local_vals=local_vals)
 
 
 def assert_eigen_shifts_is_consistent_with_eigen_components(
@@ -111,15 +109,13 @@ def _unitary_(self):
             val_from_trace, val_from_unitary)
 
 
-def _assert_meets_standards_helper(
-        val: Any,
-        qubit_count: Optional[int],
-        ignoring_global_phase,
-        setup_code: str,
-        global_vals: Optional[Dict[str, Any]],
-        local_vals: Optional[Dict[str, Any]]) -> None:
-    assert_has_consistent_qid_shape(val, qubit_count=qubit_count)
-    assert_has_consistent_apply_unitary(val, qubit_count=qubit_count)
+def _assert_meets_standards_helper(val: Any, *, ignoring_global_phase: bool,
+                                   setup_code: str,
+                                   global_vals: Optional[Dict[str, Any]],
+                                   local_vals: Optional[Dict[str, Any]]
+                                  ) -> None:
+    assert_has_consistent_qid_shape(val)
+    assert_has_consistent_apply_unitary(val)
     assert_qasm_is_consistent_with_unitary(val)
     assert_has_consistent_trace_distance_bound(val)
     assert_decompose_is_consistent_with_unitary(val,
diff --git a/cirq/testing/equivalent_repr_eval.py b/cirq/testing/equivalent_repr_eval.py
--- a/cirq/testing/equivalent_repr_eval.py
+++ b/cirq/testing/equivalent_repr_eval.py
@@ -28,6 +28,10 @@ def assert_equivalent_repr(
             code that produces an equivalent value.
         setup_code: Code that must be executed before the repr can be evaluated.
             Ideally this should just be a series of 'import' lines.
+        global_vals: Pre-defined values that should be in the global scope when
+            evaluating the repr.
+        local_vals: Pre-defined values that should be in the local scope when
+            evaluating the repr.
     """
     global_vals = global_vals or {}
     local_vals = local_vals or {}