Repository: quantumlib/Cirq

Issue:
WaitGate on multiple qubits gets confused about num_qubits when resolved
**Description of the issue**
`WaitGate.num_qubits()` appears to mutate when going through parameter resolution.

**How to reproduce the issue**

```python
g = cirq.wait(cirq.LineQubit(1), cirq.LineQubit(2))
cirq.resolve_parameters(g, {'a': 1})
```

<details>
File ~/.virtualenvs/pyle/lib/python3.8/site-packages/cirq/protocols/resolve_parameters.py:184, in resolve_parameters(val, param_resolver, recursive)
    182     result = NotImplemented
    183 else:
--> 184     result = getter(param_resolver, recursive)
    186 if result is not NotImplemented:
    187     return result

File ~/.virtualenvs/pyle/lib/python3.8/site-packages/cirq/ops/gate_operation.py:278, in GateOperation._resolve_parameters_(self, resolver, recursive)
    274 def _resolve_parameters_(
    275     self, resolver: 'cirq.ParamResolver', recursive: bool
    276 ) -> 'cirq.Operation':
    277     resolved_gate = protocols.resolve_parameters(self.gate, resolver, recursive)
--> 278     return self.with_gate(resolved_gate)

File ~/.virtualenvs/pyle/lib/python3.8/site-packages/cirq/ops/gate_operation.py:82, in GateOperation.with_gate(self, new_gate)
     79 if self.gate is new_gate:
     80     # As GateOperation is immutable, this can return the original.
     81     return self
---> 82 return new_gate.on(*self.qubits)

File ~/.virtualenvs/pyle/lib/python3.8/site-packages/cirq/ops/raw_types.py:219, in Gate.on(self, *qubits)
    213 def on(self, *qubits: Qid) -> 'Operation':
    214     """Returns an application of this gate to the given qubits.
    215 
    216     Args:
    217         *qubits: The collection of qubits to potentially apply the gate to.
    218     """
--> 219     return ops.gate_operation.GateOperation(self, list(qubits))

File ~/.virtualenvs/pyle/lib/python3.8/site-packages/cirq/ops/gate_operation.py:61, in GateOperation.__init__(self, gate, qubits)
     54 def __init__(self, gate: 'cirq.Gate', qubits: Sequence['cirq.Qid']) -> None:
     55     """Inits GateOperation.
     56 
     57     Args:
     58         gate: The gate to apply.
     59         qubits: The qubits to operate on.
     60     """
---> 61     gate.validate_args(qubits)
     62     self._gate = gate
     63     self._qubits = tuple(qubits)

File ~/.virtualenvs/pyle/lib/python3.8/site-packages/cirq/ops/raw_types.py:211, in Gate.validate_args(self, qubits)
    194 def validate_args(self, qubits: Sequence['cirq.Qid']) -> None:
    195     """Checks if this gate can be applied to the given qubits.
    196 
    197     By default checks that:
   (...)
    209         ValueError: The gate can't be applied to the qubits.
    210     """
--> 211     _validate_qid_shape(self, qubits)

File ~/.virtualenvs/pyle/lib/python3.8/site-packages/cirq/ops/raw_types.py:929, in _validate_qid_shape(val, qubits)
    927 qid_shape = protocols.qid_shape(val)
    928 if len(qubits) != len(qid_shape):
--> 929     raise ValueError(
    930         'Wrong number of qubits for <{!r}>. '
    931         'Expected {} qubits but got <{!r}>.'.format(val, len(qid_shape), qubits)
    932     )
    933 if any(qid.dimension != dimension for qid, dimension in zip(qubits, qid_shape)):
    934     raise ValueError(
    935         'Wrong shape of qids for <{!r}>. '
    936         'Expected {} but got {} <{!r}>.'.format(
    937             val, qid_shape, tuple(qid.dimension for qid in qubits), qubits
    938         )
    939     )

ValueError: Wrong number of qubits for <cirq.WaitGate(cirq.Duration(millis=0))>. Expected 1 qubits but got <[cirq.LineQubit(1), cirq.LineQubit(2)]>.
</details>

**Cirq version**
'0.15.0.dev20220503224557'





Discussion:

Resolving patch:
diff --git a/cirq-core/cirq/ops/wait_gate.py b/cirq-core/cirq/ops/wait_gate.py
--- a/cirq-core/cirq/ops/wait_gate.py
+++ b/cirq-core/cirq/ops/wait_gate.py
@@ -78,7 +78,10 @@ def _parameter_names_(self) -> AbstractSet[str]:
         return protocols.parameter_names(self.duration)
 
     def _resolve_parameters_(self, resolver: 'cirq.ParamResolver', recursive: bool) -> 'WaitGate':
-        return WaitGate(protocols.resolve_parameters(self.duration, resolver, recursive))
+        return WaitGate(
+            protocols.resolve_parameters(self.duration, resolver, recursive),
+            qid_shape=self._qid_shape,
+        )
 
     def _qid_shape_(self) -> Tuple[int, ...]:
         return self._qid_shape


Test patch:
diff --git a/cirq-core/cirq/ops/wait_gate_test.py b/cirq-core/cirq/ops/wait_gate_test.py
--- a/cirq-core/cirq/ops/wait_gate_test.py
+++ b/cirq-core/cirq/ops/wait_gate_test.py
@@ -76,6 +76,14 @@ def test_qid_shape():
         cirq.WaitGate(0, qid_shape=(2, 2), num_qubits=1)
 
 
+@pytest.mark.parametrize('num_qubits', [1, 2, 3])
+def test_resolve_parameters(num_qubits: int) -> None:
+    gate = cirq.WaitGate(duration=cirq.Duration(nanos=sympy.Symbol('t_ns')), num_qubits=num_qubits)
+    resolved = cirq.resolve_parameters(gate, {'t_ns': 10})
+    assert resolved.duration == cirq.Duration(nanos=10)
+    assert cirq.num_qubits(resolved) == num_qubits
+
+
 def test_json():
     q0, q1 = cirq.GridQubit.rect(1, 2)
     qtrit = cirq.GridQid(1, 2, dimension=3)