From becdc64d7c31d5c140b01906f52799b18839c610 Mon Sep 17 00:00:00 2001
From: Bicheng Ying <bichengying@gmail.com>
Date: Tue, 15 Jul 2025 03:59:20 +0000
Subject: [PATCH 01/10] Init code for generic analog circuit

---
 .../analog_experiments/__init__.py            |   4 +
 .../analog_trajectory_util.py                 |  23 ++-
 .../analog_trajectory_util_test.py            |  38 ++++-
 .../generic_analog_circuit.py                 | 141 ++++++++++++++++++
 .../generic_analog_circuit_test.py            | 116 ++++++++++++++
 cirq-google/cirq_google/study/symbol_util.py  |  25 ++++
 .../cirq_google/study/symbol_util_test.py     |   7 +
 7 files changed, 348 insertions(+), 6 deletions(-)
 create mode 100644 cirq-google/cirq_google/experimental/analog_experiments/generic_analog_circuit.py
 create mode 100644 cirq-google/cirq_google/experimental/analog_experiments/generic_analog_circuit_test.py

diff --git a/cirq-google/cirq_google/experimental/analog_experiments/__init__.py b/cirq-google/cirq_google/experimental/analog_experiments/__init__.py
index 13bda79ea9b..9c62b52fc87 100644
--- a/cirq-google/cirq_google/experimental/analog_experiments/__init__.py
+++ b/cirq-google/cirq_google/experimental/analog_experiments/__init__.py
@@ -18,3 +18,7 @@
     FrequencyMap as FrequencyMap,
     AnalogTrajectory as AnalogTrajectory,
 )
+
+from cirq_google.experimental.analog_experiments.generic_analog_circuit import (
+    GenericAnalogCircuit as GenericAnalogCircuit,
+)
diff --git a/cirq-google/cirq_google/experimental/analog_experiments/analog_trajectory_util.py b/cirq-google/cirq_google/experimental/analog_experiments/analog_trajectory_util.py
index 4635aa4370e..8b617be65c3 100644
--- a/cirq-google/cirq_google/experimental/analog_experiments/analog_trajectory_util.py
+++ b/cirq-google/cirq_google/experimental/analog_experiments/analog_trajectory_util.py
@@ -19,6 +19,7 @@
 import attrs
 import matplotlib.pyplot as plt
 import numpy as np
+import sympy
 import tunits as tu
 
 import cirq
@@ -36,11 +37,13 @@ class FrequencyMap:
         duration: duration of step
         qubit_freqs: dict describing qubit frequencies at end of step (None if idle)
         couplings: dict describing coupling rates at end of step
+        is_wait_step: a bool indicating only wait gate should be added.
     """
 
     duration: su.ValueOrSymbol
     qubit_freqs: dict[str, su.ValueOrSymbol | None]
     couplings: dict[tuple[str, str], su.ValueOrSymbol]
+    is_wait_step: bool
 
     def _is_parameterized_(self) -> bool:
         return (
@@ -68,8 +71,22 @@ def _resolve_parameters_(
             couplings={
                 k: su.direct_symbol_replacement(v, resolver_) for k, v in self.couplings.items()
             },
+            is_wait_step=self.is_wait_step,
         )
 
+    def duration_nanos(self) -> float | su.TunitForDurationNanos:
+        # The following is the patching code for symbol/tunit can work
+        # with cirq.Duration object.
+        if isinstance(self.duration, tu.Value):
+            nanos = self.duration[tu.ns]
+        elif isinstance(self.duration, sympy.Symbol):
+            nanos = su.TunitForDurationNanos(self.duration)
+        else:
+            raise ValueError(
+                "The duration in the freq map must either be a tu.Value or a sympy.Symbol."
+            )
+        return nanos
+
 
 class AnalogTrajectory:
     """Class for handling qubit frequency and coupling trajectories that
@@ -129,9 +146,11 @@ def from_sparse_trajectory(
         full_trajectory: list[FrequencyMap] = []
         init_qubit_freq_dict: dict[str, tu.Value | None] = {q: None for q in qubits}
         init_g_dict: dict[tuple[str, str], tu.Value] = {p: 0 * tu.MHz for p in pairs}
-        full_trajectory.append(FrequencyMap(0 * tu.ns, init_qubit_freq_dict, init_g_dict))
+        full_trajectory.append(FrequencyMap(0 * tu.ns, init_qubit_freq_dict, init_g_dict, False))
 
         for dt, qubit_freq_dict, g_dict in sparse_trajectory:
+            # When both qubit_freq_dict and g_dict is empty, it is a wait step.
+            is_wait_step = not (qubit_freq_dict or g_dict)
             # If no freq provided, set equal to previous
             new_qubit_freq_dict = {
                 q: qubit_freq_dict.get(q, full_trajectory[-1].qubit_freqs.get(q)) for q in qubits
@@ -141,7 +160,7 @@ def from_sparse_trajectory(
                 p: g_dict.get(p, full_trajectory[-1].couplings.get(p)) for p in pairs  # type: ignore[misc]
             }
 
-            full_trajectory.append(FrequencyMap(dt, new_qubit_freq_dict, new_g_dict))
+            full_trajectory.append(FrequencyMap(dt, new_qubit_freq_dict, new_g_dict, is_wait_step))
         return cls(full_trajectory=full_trajectory, qubits=qubits, pairs=pairs)
 
     def get_full_trajectory_with_resolved_idles(
diff --git a/cirq-google/cirq_google/experimental/analog_experiments/analog_trajectory_util_test.py b/cirq-google/cirq_google/experimental/analog_experiments/analog_trajectory_util_test.py
index e57dfd58411..89bf0e289af 100644
--- a/cirq-google/cirq_google/experimental/analog_experiments/analog_trajectory_util_test.py
+++ b/cirq-google/cirq_google/experimental/analog_experiments/analog_trajectory_util_test.py
@@ -26,6 +26,7 @@ def freq_map() -> atu.FrequencyMap:
         10 * tu.ns,
         {"q0_0": 5 * tu.GHz, "q0_1": 6 * tu.GHz, "q0_2": sympy.Symbol("f_q0_2")},
         {("q0_0", "q0_1"): 5 * tu.MHz, ("q0_1", "q0_2"): sympy.Symbol("g_q0_1_q0_2")},
+        False,
     )
 
 
@@ -42,9 +43,17 @@ def test_freq_map_resolve(freq_map: atu.FrequencyMap) -> None:
         10 * tu.ns,
         {"q0_0": 5 * tu.GHz, "q0_1": 6 * tu.GHz, "q0_2": 6 * tu.GHz},
         {("q0_0", "q0_1"): 5 * tu.MHz, ("q0_1", "q0_2"): 7 * tu.MHz},
+        False,
     )
 
 
+def test_duration_nanos_in_freq_map() -> None:
+    x = atu.FrequencyMap(10 * tu.ns, {}, {}, False)
+    assert x.duration_nanos() == 10
+    x = atu.FrequencyMap(sympy.Symbol("t"), {}, {}, False)
+    assert cirq.resolve_parameters(x.duration_nanos(), {"t": 10 * tu.ns}) == 10
+
+
 FreqMapType = tuple[tu.Value, dict[str, tu.Value | None], dict[tuple[str, str], tu.Value]]
 
 
@@ -52,36 +61,47 @@ def test_freq_map_resolve(freq_map: atu.FrequencyMap) -> None:
 def sparse_trajectory() -> list[FreqMapType]:
     traj1: FreqMapType = (20 * tu.ns, {"q0_1": 5 * tu.GHz}, {})
     traj2: FreqMapType = (30 * tu.ns, {"q0_2": 8 * tu.GHz}, {})
-    traj3: FreqMapType = (
+    traj3: FreqMapType = (35 * tu.ns, {}, {})
+    traj4: FreqMapType = (
         40 * tu.ns,
         {"q0_0": 8 * tu.GHz, "q0_1": None, "q0_2": None},
         {("q0_0", "q0_1"): 5 * tu.MHz, ("q0_1", "q0_2"): 8 * tu.MHz},
     )
-    return [traj1, traj2, traj3]
+    return [traj1, traj2, traj3, traj4]
 
 
 def test_full_traj(sparse_trajectory: list[FreqMapType]) -> None:
     analog_traj = atu.AnalogTrajectory.from_sparse_trajectory(sparse_trajectory)
-    assert len(analog_traj.full_trajectory) == 4
+    assert len(analog_traj.full_trajectory) == 5
     assert analog_traj.full_trajectory[0] == atu.FrequencyMap(
         0 * tu.ns,
         {"q0_0": None, "q0_1": None, "q0_2": None},
         {("q0_0", "q0_1"): 0 * tu.MHz, ("q0_1", "q0_2"): 0 * tu.MHz},
+        False,
     )
     assert analog_traj.full_trajectory[1] == atu.FrequencyMap(
         20 * tu.ns,
         {"q0_0": None, "q0_1": 5 * tu.GHz, "q0_2": None},
         {("q0_0", "q0_1"): 0 * tu.MHz, ("q0_1", "q0_2"): 0 * tu.MHz},
+        False,
     )
     assert analog_traj.full_trajectory[2] == atu.FrequencyMap(
         30 * tu.ns,
         {"q0_0": None, "q0_1": 5 * tu.GHz, "q0_2": 8 * tu.GHz},
         {("q0_0", "q0_1"): 0 * tu.MHz, ("q0_1", "q0_2"): 0 * tu.MHz},
+        False,
     )
     assert analog_traj.full_trajectory[3] == atu.FrequencyMap(
+        35 * tu.ns,
+        {"q0_0": None, "q0_1": 5 * tu.GHz, "q0_2": 8 * tu.GHz},
+        {("q0_0", "q0_1"): 0 * tu.MHz, ("q0_1", "q0_2"): 0 * tu.MHz},
+        True,
+    )
+    assert analog_traj.full_trajectory[4] == atu.FrequencyMap(
         40 * tu.ns,
         {"q0_0": 8 * tu.GHz, "q0_1": None, "q0_2": None},
         {("q0_0", "q0_1"): 5 * tu.MHz, ("q0_1", "q0_2"): 8 * tu.MHz},
+        False,
     )
 
 
@@ -92,26 +112,36 @@ def test_get_full_trajectory_with_resolved_idles(sparse_trajectory: list[FreqMap
         {"q0_0": 5 * tu.GHz, "q0_1": 6 * tu.GHz, "q0_2": 7 * tu.GHz}
     )
 
-    assert len(resolved_full_traj) == 4
+    assert len(resolved_full_traj) == 5
     assert resolved_full_traj[0] == atu.FrequencyMap(
         0 * tu.ns,
         {"q0_0": 5 * tu.GHz, "q0_1": 6 * tu.GHz, "q0_2": 7 * tu.GHz},
         {("q0_0", "q0_1"): 0 * tu.MHz, ("q0_1", "q0_2"): 0 * tu.MHz},
+        False,
     )
     assert resolved_full_traj[1] == atu.FrequencyMap(
         20 * tu.ns,
         {"q0_0": 5 * tu.GHz, "q0_1": 5 * tu.GHz, "q0_2": 7 * tu.GHz},
         {("q0_0", "q0_1"): 0 * tu.MHz, ("q0_1", "q0_2"): 0 * tu.MHz},
+        False,
     )
     assert resolved_full_traj[2] == atu.FrequencyMap(
         30 * tu.ns,
         {"q0_0": 5 * tu.GHz, "q0_1": 5 * tu.GHz, "q0_2": 8 * tu.GHz},
         {("q0_0", "q0_1"): 0 * tu.MHz, ("q0_1", "q0_2"): 0 * tu.MHz},
+        False,
     )
     assert resolved_full_traj[3] == atu.FrequencyMap(
+        35 * tu.ns,
+        {"q0_0": 5 * tu.GHz, "q0_1": 5 * tu.GHz, "q0_2": 8 * tu.GHz},
+        {("q0_0", "q0_1"): 0 * tu.MHz, ("q0_1", "q0_2"): 0 * tu.MHz},
+        True,
+    )
+    assert resolved_full_traj[4] == atu.FrequencyMap(
         40 * tu.ns,
         {"q0_0": 8 * tu.GHz, "q0_1": 6 * tu.GHz, "q0_2": 7 * tu.GHz},
         {("q0_0", "q0_1"): 5 * tu.MHz, ("q0_1", "q0_2"): 8 * tu.MHz},
+        False,
     )
 
 
diff --git a/cirq-google/cirq_google/experimental/analog_experiments/generic_analog_circuit.py b/cirq-google/cirq_google/experimental/analog_experiments/generic_analog_circuit.py
new file mode 100644
index 00000000000..bedd3bc1eab
--- /dev/null
+++ b/cirq-google/cirq_google/experimental/analog_experiments/generic_analog_circuit.py
@@ -0,0 +1,141 @@
+# Copyright 2025 The Cirq Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+import functools
+import re
+
+import cirq
+from cirq_google.experimental.analog_experiments import analog_trajectory_util as atu
+from cirq_google.ops import analog_detune_gates as adg
+from cirq_google.study import symbol_util as su
+
+
+def _get_neighbor_freqs(
+    qubit_pair: tuple[str, str], qubit_freq_dict: dict[str, su.ValueOrSymbol | None]
+) -> tuple[su.ValueOrSymbol | None, su.ValueOrSymbol | None]:
+    """Get neighbor freqs from qubit_freq_dict given the pair."""
+    sorted_pair = sorted(qubit_pair, key=_to_grid_qubit)
+    return (qubit_freq_dict[sorted_pair[0]], qubit_freq_dict[sorted_pair[1]])
+
+
+@functools.cache
+def _to_grid_qubit(qubit_name: str) -> cirq.GridQubit:
+    match = re.compile(r"^q(\d+)_(\d+)$").match(qubit_name)
+    return cirq.GridQubit(int(match[1]), int(match[2]))
+
+
+def _coupler_name_from_qubit_pair(qubit_pair: tuple[str, str]) -> str:
+    sorted_pair = sorted(qubit_pair, key=_to_grid_qubit)
+    return f"c_{sorted_pair[0]}_{sorted_pair[1]}"
+
+
+def _get_neighbor_coupler_freqs(
+    qubit_name: str, coupler_g_dict: dict[tuple[str, str], su.ValueOrSymbol]
+) -> dict[str, su.ValueOrSymbol]:
+    """Get neighbor coupler coupling strength g given qubit name."""
+    return {
+        _coupler_name_from_qubit_pair(pair): g
+        for pair, g in coupler_g_dict.items()
+        if qubit_name in pair
+    }
+
+
+class GenericAnalogCircuit:
+    """Class for making arbitrary analog circuits. The circuit is defined by an
+    AnalogTrajectory object. The class constructs the circuit from AnalogDetune
+    pulses, which automatically calculate the necessary bias amps to both qubits
+    and couplers, using tu.Values from analog calibration whenever available.
+
+    Attributes:
+        trajectory: AnalogTrajectory object defining the circuit
+        g_ramp_shaping: coupling ramps are shaped according to ramp_shape_exp if True
+        qubits: list of qubits in the circuit
+        pairs: list of couplers in the circuit
+        ramp_shape_exp: exponent of g_ramp (g proportional to t^ramp_shape_exp)
+        interpolate_coupling_cal: interpolates between calibrated coupling tu.Values if True
+        linear_qubit_ramp: if True, the qubit ramp is linear. if false, a cosine shaped
+            ramp is used.
+    """
+
+    def __init__(
+        self,
+        trajectory: atu.AnalogTrajectory,
+        g_ramp_shaping: bool = False,
+        ramp_shape_exp: int = 1,
+        interpolate_coupling_cal: bool = False,
+        linear_qubit_ramp: bool = True,
+    ):
+        self.trajectory = trajectory
+        self.g_ramp_shaping = g_ramp_shaping
+        self.ramp_shape_exp = ramp_shape_exp
+        self.interpolate_coupling_cal = interpolate_coupling_cal
+        self.linear_qubit_ramp = linear_qubit_ramp
+
+    def make_circuit(self) -> cirq.Circuit:
+        """Assemble moments described in trajectory."""
+        prev_freq_map = self.trajectory.full_trajectory[0]
+        moments = []
+        for freq_map in self.trajectory.full_trajectory[1:]:
+            if freq_map.is_wait_step:
+                # `duration_nanos` is a patch solution for tunit working with cirq.Duration.
+                moment = cirq.Moment(
+                    cirq.wait(self.trajectory.qubits, nanos=freq_map.duration_nanos())  # type: ignore
+                )
+            else:
+                moment = self.make_one_moment(freq_map, prev_freq_map)
+            moments.append(moment)
+            prev_freq_map = freq_map
+
+        return cirq.Circuit.from_moments(*moments)
+
+    def make_one_moment(
+        self, freq_map: atu.FrequencyMap, prev_freq_map: atu.FrequencyMap
+    ) -> cirq.Moment:
+        """Make one moment of analog detune qubit and coupler gates given freqs."""
+        qubit_gates = []
+        for q, freq in freq_map.qubit_freqs.items():
+            qubit_gates.append(
+                adg.AnalogDetuneQubit(
+                    length=freq_map.duration,
+                    w=freq_map.duration,
+                    target_freq=freq,
+                    prev_freq=prev_freq_map.qubit_freqs.get(q),
+                    neighbor_coupler_g_dict=_get_neighbor_coupler_freqs(q, freq_map.couplings),
+                    prev_neighbor_coupler_g_dict=_get_neighbor_coupler_freqs(
+                        q, prev_freq_map.couplings
+                    ),
+                    linear_rise=self.linear_qubit_ramp,
+                ).on(_to_grid_qubit(q))
+            )
+        coupler_gates = []
+        for p, g_max in freq_map.couplings.items():
+            # Currently skipping the step if these are the same.
+            # However, change in neighbor qubit freq could potentially change coupler amp
+            if g_max == prev_freq_map.couplings[p]:
+                continue
+
+            coupler_gates.append(
+                adg.AnalogDetuneCouplerOnly(
+                    length=freq_map.duration,
+                    w=freq_map.duration,
+                    g_0=prev_freq_map.couplings[p],
+                    g_max=g_max,
+                    g_ramp_exponent=self.ramp_shape_exp,
+                    neighbor_qubits_freq=_get_neighbor_freqs(p, freq_map.qubit_freqs),
+                    prev_neighbor_qubits_freq=_get_neighbor_freqs(p, prev_freq_map.qubit_freqs),
+                    interpolate_coupling_cal=self.interpolate_coupling_cal,
+                ).on(*sorted([_to_grid_qubit(p[0]), _to_grid_qubit(p[1])]))
+            )
+
+        return cirq.Moment(qubit_gates + coupler_gates)
diff --git a/cirq-google/cirq_google/experimental/analog_experiments/generic_analog_circuit_test.py b/cirq-google/cirq_google/experimental/analog_experiments/generic_analog_circuit_test.py
new file mode 100644
index 00000000000..307440e3666
--- /dev/null
+++ b/cirq-google/cirq_google/experimental/analog_experiments/generic_analog_circuit_test.py
@@ -0,0 +1,116 @@
+# Copyright 2025 The Cirq Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+import sympy
+import tunits as tu
+
+import cirq
+from cirq_google.experimental.analog_experiments import (
+    analog_trajectory_util as atu,
+    generic_analog_circuit as gac,
+)
+from cirq_google.ops.analog_detune_gates import AnalogDetuneCouplerOnly, AnalogDetuneQubit
+
+
+def test_get_neighbor_freqs():
+    pair = ("q0_0", "q0_1")
+    qubit_freq_dict = {"q0_0": 5 * tu.GHz, "q0_1": sympy.Symbol("f_q"), "q0_2": 6 * tu.GHz}
+    neighbor_freqs = gac._get_neighbor_freqs(pair, qubit_freq_dict)
+    assert neighbor_freqs == (5 * tu.GHz, sympy.Symbol("f_q"))
+
+
+def test_coupler_name_from_qubit_pair():
+    pair = ("q0_0", "q0_1")
+    coupler_name = gac._coupler_name_from_qubit_pair(pair)
+    assert coupler_name == "c_q0_0_q0_1"
+
+    pair = ("q9_0", "q10_0")
+    coupler_name = gac._coupler_name_from_qubit_pair(pair)
+    assert coupler_name == "c_q9_0_q10_0"
+
+    pair = ("q7_8", "q7_7")
+    coupler_name = gac._coupler_name_from_qubit_pair(pair)
+    assert coupler_name == "c_q7_7_q7_8"
+
+
+def test_make_one_moment_of_generic_analog_circuit():
+    trajectory = None  # we don't need trajector in this test.
+    generic_analog_circuit = gac.GenericAnalogCircuit(trajectory)
+
+    freq_map = atu.FrequencyMap(
+        duration=3 * tu.ns,
+        qubit_freqs={"q0_0": 5 * tu.GHz, "q0_1": 6 * tu.GHz, "q0_2": sympy.Symbol("f_q0_2")},
+        couplings={("q0_0", "q0_1"): 5 * tu.MHz, ("q0_1", "q0_2"): 6 * tu.MHz},
+        is_wait_step=False,
+    )
+    prev_freq_map = atu.FrequencyMap(
+        duration=9 * tu.ns,
+        qubit_freqs={"q0_0": 4 * tu.GHz, "q0_1": 6 * tu.GHz, "q0_2": sympy.Symbol("f_q0_2")},
+        couplings={("q0_0", "q0_1"): 2 * tu.MHz, ("q0_1", "q0_2"): 3 * tu.MHz},
+        is_wait_step=False,
+    )
+
+    moment = generic_analog_circuit.make_one_moment(freq_map, prev_freq_map)
+
+    assert len(moment.operations) == 5
+    # Three detune qubit gates
+    assert moment.operations[0] == AnalogDetuneQubit(
+        length=3 * tu.ns,
+        w=3 * tu.ns,
+        target_freq=5 * tu.GHz,
+        prev_freq=4 * tu.GHz,
+        neighbor_coupler_g_dict={"c_q0_0_q0_1": 5 * tu.MHz},
+        prev_neighbor_coupler_g_dict={"c_q0_0_q0_1": 2 * tu.MHz},
+        linear_rise=True,
+    ).on(cirq.GridQubit(0, 0))
+    assert moment.operations[1] == AnalogDetuneQubit(
+        length=3 * tu.ns,
+        w=3 * tu.ns,
+        target_freq=6 * tu.GHz,
+        prev_freq=6 * tu.GHz,
+        neighbor_coupler_g_dict={"c_q0_0_q0_1": 5 * tu.MHz, "c_q0_1_q0_2": 6 * tu.MHz},
+        prev_neighbor_coupler_g_dict={"c_q0_0_q0_1": 2 * tu.MHz, "c_q0_1_q0_2": 3 * tu.MHz},
+        linear_rise=True,
+    ).on(cirq.GridQubit(0, 1))
+    assert moment.operations[2] == AnalogDetuneQubit(
+        length=3 * tu.ns,
+        w=3 * tu.ns,
+        target_freq=sympy.Symbol("f_q0_2"),
+        prev_freq=sympy.Symbol("f_q0_2"),
+        neighbor_coupler_g_dict={"c_q0_1_q0_2": 6 * tu.MHz},
+        prev_neighbor_coupler_g_dict={"c_q0_1_q0_2": 3 * tu.MHz},
+        linear_rise=True,
+    ).on(cirq.GridQubit(0, 2))
+
+    # Two detune coupler only gates
+    assert moment.operations[3] == AnalogDetuneCouplerOnly(
+        length=3 * tu.ns,
+        w=3 * tu.ns,
+        g_0=2 * tu.MHz,
+        g_max=5 * tu.MHz,
+        g_ramp_exponent=1,
+        neighbor_qubits_freq=(5 * tu.GHz, 6 * tu.GHz),
+        prev_neighbor_qubits_freq=(4 * tu.GHz, 6 * tu.GHz),
+        interpolate_coupling_cal=False,
+    ).on(cirq.GridQubit(0, 0), cirq.GridQubit(0, 1))
+    assert moment.operations[4] == AnalogDetuneCouplerOnly(
+        length=3 * tu.ns,
+        w=3 * tu.ns,
+        g_0=3 * tu.MHz,
+        g_max=6 * tu.MHz,
+        g_ramp_exponent=1,
+        neighbor_qubits_freq=(6 * tu.GHz, sympy.Symbol("f_q0_2")),
+        prev_neighbor_qubits_freq=(6 * tu.GHz, sympy.Symbol("f_q0_2")),
+        interpolate_coupling_cal=False,
+    ).on(cirq.GridQubit(0, 1), cirq.GridQubit(0, 2))
diff --git a/cirq-google/cirq_google/study/symbol_util.py b/cirq-google/cirq_google/study/symbol_util.py
index b2360543201..86142e11b4d 100644
--- a/cirq-google/cirq_google/study/symbol_util.py
+++ b/cirq-google/cirq_google/study/symbol_util.py
@@ -51,3 +51,28 @@ def is_parameterized_dict(dict_with_value: dict[Any, ValueOrSymbol] | None) -> b
     if dict_with_value is None:
         return False  # pragma: no cover
     return any(cirq.is_parameterized(v) for v in dict_with_value.values())
+
+
+class TunitForDurationNanos:
+    """A wrapper class that can be used with symbols for duration nanos.
+
+    When resolving it, it will nanos (as float) so that we can use it
+    as the input of `cirq.Duration(nanos=TunitForDurationNanos())`.
+    """
+
+    def __init__(self, duration: sympy.Symbol):
+        if not isinstance(duration, sympy.Symbol):
+            raise ValueError("The duration must be a symbol.")
+        self.duration = duration
+
+    def _is_parameterized_(self) -> bool:
+        return cirq.is_parameterized(self.duration)
+
+    def _parameter_names_(self) -> AbstractSet[str]:
+        return cirq.parameter_names(self.duration)
+
+    def _resolve_parameters_(
+        self, resolver: cirq.ParamResolverOrSimilarType, recursive: bool
+    ) -> float:
+        resolver_ = cirq.ParamResolver(resolver)
+        return direct_symbol_replacement(self.duration, resolver_)[tu.ns]
diff --git a/cirq-google/cirq_google/study/symbol_util_test.py b/cirq-google/cirq_google/study/symbol_util_test.py
index a6ef9d50b4d..4efae688e28 100644
--- a/cirq-google/cirq_google/study/symbol_util_test.py
+++ b/cirq-google/cirq_google/study/symbol_util_test.py
@@ -46,3 +46,10 @@ def test_direct_symbol_replacement():
     value_resolved = [su.direct_symbol_replacement(v, resolver) for v in value_list]
 
     assert value_resolved == [5 * tu.ns, 8 * tu.GHz, sympy.Symbol("z"), 123, "fd"]
+
+
+def test_tunit_for_duration_nanos():
+    xx = su.TunitForDurationNanos(duration=sympy.Symbol("t"))
+    assert cirq.is_parameterized(xx)
+    assert xx._parameter_names_() == {"t"}
+    assert cirq.resolve_parameters(xx, {"t": 5 * tu.ns}) == 5

From 94f48c78fd86cc12ed0e472c96781db77a3d8e86 Mon Sep 17 00:00:00 2001
From: Bicheng Ying <bichengying@gmail.com>
Date: Tue, 15 Jul 2025 18:36:01 +0000
Subject: [PATCH 02/10] Add test and refine the TunitForDurationNanos

---
 .../analog_experiments/__init__.py            |  2 +-
 .../analog_trajectory_util_test.py            | 15 +++++---
 .../generic_analog_circuit.py                 | 21 ++++++++---
 .../generic_analog_circuit_test.py            | 37 +++++++++++++++++--
 cirq-google/cirq_google/study/symbol_util.py  | 27 ++++++++++----
 .../cirq_google/study/symbol_util_test.py     | 19 ++++++++--
 6 files changed, 94 insertions(+), 27 deletions(-)

diff --git a/cirq-google/cirq_google/experimental/analog_experiments/__init__.py b/cirq-google/cirq_google/experimental/analog_experiments/__init__.py
index 9c62b52fc87..63af6311a74 100644
--- a/cirq-google/cirq_google/experimental/analog_experiments/__init__.py
+++ b/cirq-google/cirq_google/experimental/analog_experiments/__init__.py
@@ -20,5 +20,5 @@
 )
 
 from cirq_google.experimental.analog_experiments.generic_analog_circuit import (
-    GenericAnalogCircuit as GenericAnalogCircuit,
+    GenericAnalogCircuitBuilder as GenericAnalogCircuitBuilder,
 )
diff --git a/cirq-google/cirq_google/experimental/analog_experiments/analog_trajectory_util_test.py b/cirq-google/cirq_google/experimental/analog_experiments/analog_trajectory_util_test.py
index 89bf0e289af..64a2e39ffc9 100644
--- a/cirq-google/cirq_google/experimental/analog_experiments/analog_trajectory_util_test.py
+++ b/cirq-google/cirq_google/experimental/analog_experiments/analog_trajectory_util_test.py
@@ -48,10 +48,13 @@ def test_freq_map_resolve(freq_map: atu.FrequencyMap) -> None:
 
 
 def test_duration_nanos_in_freq_map() -> None:
-    x = atu.FrequencyMap(10 * tu.ns, {}, {}, False)
-    assert x.duration_nanos() == 10
-    x = atu.FrequencyMap(sympy.Symbol("t"), {}, {}, False)
-    assert cirq.resolve_parameters(x.duration_nanos(), {"t": 10 * tu.ns}) == 10
+    fm = atu.FrequencyMap(10 * tu.ns, {}, {}, False)
+    assert fm.duration_nanos() == 10
+    fm = atu.FrequencyMap(sympy.Symbol("t"), {}, {}, False)
+    assert cirq.resolve_parameters(fm.duration_nanos(), {"t": 10 * tu.ns}).duration_nanos == 10
+
+    with pytest.raises(ValueError, match="either be a tu.Value or a sympy.Symbol"):
+        atu.FrequencyMap(10, {}, {}, False).duration_nanos()
 
 
 FreqMapType = tuple[tu.Value, dict[str, tu.Value | None], dict[tuple[str, str], tu.Value]]
@@ -145,7 +148,7 @@ def test_get_full_trajectory_with_resolved_idles(sparse_trajectory: list[FreqMap
     )
 
 
-def test_plot_with_unresolved_parameters():
+def test_plot_with_unresolved_parameters() -> None:
     traj1: FreqMapType = (20 * tu.ns, {"q0_1": sympy.Symbol("qf")}, {})
     traj2: FreqMapType = (sympy.Symbol("t"), {"q0_2": 8 * tu.GHz}, {})
     analog_traj = atu.AnalogTrajectory.from_sparse_trajectory([traj1, traj2])
@@ -154,7 +157,7 @@ def test_plot_with_unresolved_parameters():
         analog_traj.plot()
 
 
-def test_analog_traj_plot():
+def test_analog_traj_plot() -> None:
     traj1: FreqMapType = (5 * tu.ns, {"q0_1": sympy.Symbol("qf")}, {("q0_0", "q0_1"): 2 * tu.MHz})
     traj2: FreqMapType = (sympy.Symbol("t"), {"q0_2": 8 * tu.GHz}, {})
     analog_traj = atu.AnalogTrajectory.from_sparse_trajectory([traj1, traj2])
diff --git a/cirq-google/cirq_google/experimental/analog_experiments/generic_analog_circuit.py b/cirq-google/cirq_google/experimental/analog_experiments/generic_analog_circuit.py
index bedd3bc1eab..d7ef0af0c6e 100644
--- a/cirq-google/cirq_google/experimental/analog_experiments/generic_analog_circuit.py
+++ b/cirq-google/cirq_google/experimental/analog_experiments/generic_analog_circuit.py
@@ -51,7 +51,7 @@ def _get_neighbor_coupler_freqs(
     }
 
 
-class GenericAnalogCircuit:
+class GenericAnalogCircuitBuilder:
     """Class for making arbitrary analog circuits. The circuit is defined by an
     AnalogTrajectory object. The class constructs the circuit from AnalogDetune
     pulses, which automatically calculate the necessary bias amps to both qubits
@@ -88,10 +88,21 @@ def make_circuit(self) -> cirq.Circuit:
         moments = []
         for freq_map in self.trajectory.full_trajectory[1:]:
             if freq_map.is_wait_step:
-                # `duration_nanos` is a patch solution for tunit working with cirq.Duration.
-                moment = cirq.Moment(
-                    cirq.wait(self.trajectory.qubits, nanos=freq_map.duration_nanos())  # type: ignore
-                )
+                target = [_to_grid_qubit(q) for q in self.trajectory.qubits]
+                d = freq_map.duration_nanos()
+                if isinstance(d, float):
+                    wait_gate = cirq.WaitGate(
+                        cirq.Duration(nanos=d), qid_shape=cirq.qid_shape(target)
+                    )
+                else:
+                    # The following is patching solution for resolving the parameter
+                    # can be tunits. It should only work for pyle internal translation.
+                    wait_gate = cirq.WaitGate(
+                        cirq.Duration(nanos=1), qid_shape=cirq.qid_shape(target)
+                    )
+                    wait_gate._duration = d
+
+                moment = cirq.Moment(wait_gate.on(*target))
             else:
                 moment = self.make_one_moment(freq_map, prev_freq_map)
             moments.append(moment)
diff --git a/cirq-google/cirq_google/experimental/analog_experiments/generic_analog_circuit_test.py b/cirq-google/cirq_google/experimental/analog_experiments/generic_analog_circuit_test.py
index 307440e3666..14ee152a44a 100644
--- a/cirq-google/cirq_google/experimental/analog_experiments/generic_analog_circuit_test.py
+++ b/cirq-google/cirq_google/experimental/analog_experiments/generic_analog_circuit_test.py
@@ -45,9 +45,6 @@ def test_coupler_name_from_qubit_pair():
 
 
 def test_make_one_moment_of_generic_analog_circuit():
-    trajectory = None  # we don't need trajector in this test.
-    generic_analog_circuit = gac.GenericAnalogCircuit(trajectory)
-
     freq_map = atu.FrequencyMap(
         duration=3 * tu.ns,
         qubit_freqs={"q0_0": 5 * tu.GHz, "q0_1": 6 * tu.GHz, "q0_2": sympy.Symbol("f_q0_2")},
@@ -61,7 +58,9 @@ def test_make_one_moment_of_generic_analog_circuit():
         is_wait_step=False,
     )
 
-    moment = generic_analog_circuit.make_one_moment(freq_map, prev_freq_map)
+    trajectory = None  # we don't need trajector in this test.
+    builder = gac.GenericAnalogCircuitBuilder(trajectory)
+    moment = builder.make_one_moment(freq_map, prev_freq_map)
 
     assert len(moment.operations) == 5
     # Three detune qubit gates
@@ -114,3 +113,33 @@ def test_make_one_moment_of_generic_analog_circuit():
         prev_neighbor_qubits_freq=(6 * tu.GHz, sympy.Symbol("f_q0_2")),
         interpolate_coupling_cal=False,
     ).on(cirq.GridQubit(0, 1), cirq.GridQubit(0, 2))
+
+
+def test_generic_analog_make_circuit():
+    trajectory = atu.AnalogTrajectory.from_sparse_trajectory(
+        [
+            (5 * tu.ns, {"q0_0": 5 * tu.GHz}, {}),
+            (sympy.Symbol('t'), {}, {}),
+            (
+                10 * tu.ns,
+                {"q0_0": 8 * tu.GHz, "q0_1": sympy.Symbol('f')},
+                {("q0_0", "q0_1"): -5 * tu.MHz},
+            ),
+            (2 * tu.ns, {"q0_1": 4 * tu.GHz}, {}),
+        ]
+    )
+    builder = gac.GenericAnalogCircuitBuilder(trajectory)
+    circuit = builder.make_circuit()
+
+    assert len(circuit) == 4
+    for op in circuit[0].operations:
+        assert isinstance(op.gate, AnalogDetuneQubit)
+    for op in circuit[1].operations:
+        assert isinstance(op.gate, cirq.WaitGate)
+
+    assert isinstance(circuit[2].operations[0].gate, AnalogDetuneQubit)
+    assert isinstance(circuit[2].operations[1].gate, AnalogDetuneQubit)
+    assert isinstance(circuit[2].operations[2].gate, AnalogDetuneCouplerOnly)
+
+    for op in circuit[3].operations:
+        assert isinstance(op.gate, AnalogDetuneQubit)
diff --git a/cirq-google/cirq_google/study/symbol_util.py b/cirq-google/cirq_google/study/symbol_util.py
index 86142e11b4d..9d13d461120 100644
--- a/cirq-google/cirq_google/study/symbol_util.py
+++ b/cirq-google/cirq_google/study/symbol_util.py
@@ -60,19 +60,32 @@ class TunitForDurationNanos:
     as the input of `cirq.Duration(nanos=TunitForDurationNanos())`.
     """
 
-    def __init__(self, duration: sympy.Symbol):
-        if not isinstance(duration, sympy.Symbol):
-            raise ValueError("The duration must be a symbol.")
-        self.duration = duration
+    def __init__(self, duration_nanos: sympy.Symbol | float):
+        self.duration_nanos = duration_nanos
+
+    def total_micros(self):
+        if isinstance(self.duration_nanos, sympy.Symbol):
+            return self.duration_nanos
+        return self.duration_nanos * 1000
+
+    def total_nanos(self):
+        return self.duration_nanos
+
+    def __repr__(self) -> str:
+        if isinstance(self.duration_nanos, sympy.Symbol):
+            return f"duration={self.duration_nanos}"
+        return f"duration={self.duration_nanos} ns"
 
     def _is_parameterized_(self) -> bool:
-        return cirq.is_parameterized(self.duration)
+        return cirq.is_parameterized(self.duration_nanos)
 
     def _parameter_names_(self) -> AbstractSet[str]:
-        return cirq.parameter_names(self.duration)
+        return cirq.parameter_names(self.duration_nanos)
 
     def _resolve_parameters_(
         self, resolver: cirq.ParamResolverOrSimilarType, recursive: bool
     ) -> float:
         resolver_ = cirq.ParamResolver(resolver)
-        return direct_symbol_replacement(self.duration, resolver_)[tu.ns]
+        return TunitForDurationNanos(
+            duration_nanos=direct_symbol_replacement(self.duration_nanos, resolver_)[tu.ns]
+        )
diff --git a/cirq-google/cirq_google/study/symbol_util_test.py b/cirq-google/cirq_google/study/symbol_util_test.py
index 4efae688e28..9048237dde1 100644
--- a/cirq-google/cirq_google/study/symbol_util_test.py
+++ b/cirq-google/cirq_google/study/symbol_util_test.py
@@ -49,7 +49,18 @@ def test_direct_symbol_replacement():
 
 
 def test_tunit_for_duration_nanos():
-    xx = su.TunitForDurationNanos(duration=sympy.Symbol("t"))
-    assert cirq.is_parameterized(xx)
-    assert xx._parameter_names_() == {"t"}
-    assert cirq.resolve_parameters(xx, {"t": 5 * tu.ns}) == 5
+    d = su.TunitForDurationNanos(duration_nanos=sympy.Symbol("t"))
+    assert str(d) == "duration=t"
+    assert cirq.is_parameterized(d)
+    assert d._parameter_names_() == {"t"}
+    assert cirq.resolve_parameters(d, {"t": 5 * tu.ns}).duration_nanos == 5.0
+
+
+def test_tunit_for_duration_total_nanos():
+    d = su.TunitForDurationNanos(duration_nanos=sympy.Symbol("t"))
+    assert d.total_nanos() == sympy.Symbol("t")
+    assert d.total_micros() == sympy.Symbol("t")
+    d_resolved = cirq.resolve_parameters(d, {"t": 5 * tu.ns})
+    assert str(d_resolved) == "duration=5.0 ns"
+    assert d_resolved.total_nanos() == 5
+    assert d_resolved.total_micros() == 5000

From 9e187168d31dc629b57185f546fb9b2b9e5280ba Mon Sep 17 00:00:00 2001
From: Bicheng Ying <bichengying@gmail.com>
Date: Tue, 15 Jul 2025 18:56:29 +0000
Subject: [PATCH 03/10] test and type

---
 .../analog_trajectory_util_test.py            |  2 +-
 .../generic_analog_circuit.py                 |  4 +++-
 .../generic_analog_circuit_test.py            | 23 +++++++++++++++----
 cirq-google/cirq_google/study/symbol_util.py  |  4 +++-
 4 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/cirq-google/cirq_google/experimental/analog_experiments/analog_trajectory_util_test.py b/cirq-google/cirq_google/experimental/analog_experiments/analog_trajectory_util_test.py
index 64a2e39ffc9..eb0c01ef775 100644
--- a/cirq-google/cirq_google/experimental/analog_experiments/analog_trajectory_util_test.py
+++ b/cirq-google/cirq_google/experimental/analog_experiments/analog_trajectory_util_test.py
@@ -51,7 +51,7 @@ def test_duration_nanos_in_freq_map() -> None:
     fm = atu.FrequencyMap(10 * tu.ns, {}, {}, False)
     assert fm.duration_nanos() == 10
     fm = atu.FrequencyMap(sympy.Symbol("t"), {}, {}, False)
-    assert cirq.resolve_parameters(fm.duration_nanos(), {"t": 10 * tu.ns}).duration_nanos == 10
+    assert cirq.resolve_parameters(fm.duration_nanos(), {"t": 10 * tu.ns}).total_nanos() == 10
 
     with pytest.raises(ValueError, match="either be a tu.Value or a sympy.Symbol"):
         atu.FrequencyMap(10, {}, {}, False).duration_nanos()
diff --git a/cirq-google/cirq_google/experimental/analog_experiments/generic_analog_circuit.py b/cirq-google/cirq_google/experimental/analog_experiments/generic_analog_circuit.py
index d7ef0af0c6e..a3421b0416e 100644
--- a/cirq-google/cirq_google/experimental/analog_experiments/generic_analog_circuit.py
+++ b/cirq-google/cirq_google/experimental/analog_experiments/generic_analog_circuit.py
@@ -32,6 +32,8 @@ def _get_neighbor_freqs(
 @functools.cache
 def _to_grid_qubit(qubit_name: str) -> cirq.GridQubit:
     match = re.compile(r"^q(\d+)_(\d+)$").match(qubit_name)
+    if match is None:
+        raise ValueError(f"Invalid qubit name format: '{qubit_name}'. Expected 'q<row>_<col>'.")
     return cirq.GridQubit(int(match[1]), int(match[2]))
 
 
@@ -100,7 +102,7 @@ def make_circuit(self) -> cirq.Circuit:
                     wait_gate = cirq.WaitGate(
                         cirq.Duration(nanos=1), qid_shape=cirq.qid_shape(target)
                     )
-                    wait_gate._duration = d
+                    wait_gate._duration = d  # type: ignore
 
                 moment = cirq.Moment(wait_gate.on(*target))
             else:
diff --git a/cirq-google/cirq_google/experimental/analog_experiments/generic_analog_circuit_test.py b/cirq-google/cirq_google/experimental/analog_experiments/generic_analog_circuit_test.py
index 14ee152a44a..bd1d8ddf5b9 100644
--- a/cirq-google/cirq_google/experimental/analog_experiments/generic_analog_circuit_test.py
+++ b/cirq-google/cirq_google/experimental/analog_experiments/generic_analog_circuit_test.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import pytest
 import sympy
 import tunits as tu
 
@@ -23,14 +24,22 @@
 from cirq_google.ops.analog_detune_gates import AnalogDetuneCouplerOnly, AnalogDetuneQubit
 
 
-def test_get_neighbor_freqs():
+def test_get_neighbor_freqs() -> None:
     pair = ("q0_0", "q0_1")
     qubit_freq_dict = {"q0_0": 5 * tu.GHz, "q0_1": sympy.Symbol("f_q"), "q0_2": 6 * tu.GHz}
     neighbor_freqs = gac._get_neighbor_freqs(pair, qubit_freq_dict)
     assert neighbor_freqs == (5 * tu.GHz, sympy.Symbol("f_q"))
 
 
-def test_coupler_name_from_qubit_pair():
+def test_to_grid_qubit() -> None:
+    grid_qubit = gac._to_grid_qubit("q0_1")
+    assert grid_qubit == cirq.GridQubit(0, 1)
+
+    with pytest.raises(ValueError, match="Invalid qubit name format"):
+        gac._to_grid_qubit("q1")
+
+
+def test_coupler_name_from_qubit_pair() -> None:
     pair = ("q0_0", "q0_1")
     coupler_name = gac._coupler_name_from_qubit_pair(pair)
     assert coupler_name == "c_q0_0_q0_1"
@@ -44,7 +53,7 @@ def test_coupler_name_from_qubit_pair():
     assert coupler_name == "c_q7_7_q7_8"
 
 
-def test_make_one_moment_of_generic_analog_circuit():
+def test_make_one_moment_of_generic_analog_circuit() -> None:
     freq_map = atu.FrequencyMap(
         duration=3 * tu.ns,
         qubit_freqs={"q0_0": 5 * tu.GHz, "q0_1": 6 * tu.GHz, "q0_2": sympy.Symbol("f_q0_2")},
@@ -115,7 +124,7 @@ def test_make_one_moment_of_generic_analog_circuit():
     ).on(cirq.GridQubit(0, 1), cirq.GridQubit(0, 2))
 
 
-def test_generic_analog_make_circuit():
+def test_generic_analog_make_circuit() -> None:
     trajectory = atu.AnalogTrajectory.from_sparse_trajectory(
         [
             (5 * tu.ns, {"q0_0": 5 * tu.GHz}, {}),
@@ -125,13 +134,14 @@ def test_generic_analog_make_circuit():
                 {"q0_0": 8 * tu.GHz, "q0_1": sympy.Symbol('f')},
                 {("q0_0", "q0_1"): -5 * tu.MHz},
             ),
+            (3 * tu.ns, {}, {}),
             (2 * tu.ns, {"q0_1": 4 * tu.GHz}, {}),
         ]
     )
     builder = gac.GenericAnalogCircuitBuilder(trajectory)
     circuit = builder.make_circuit()
 
-    assert len(circuit) == 4
+    assert len(circuit) == 5
     for op in circuit[0].operations:
         assert isinstance(op.gate, AnalogDetuneQubit)
     for op in circuit[1].operations:
@@ -142,4 +152,7 @@ def test_generic_analog_make_circuit():
     assert isinstance(circuit[2].operations[2].gate, AnalogDetuneCouplerOnly)
 
     for op in circuit[3].operations:
+        assert isinstance(op.gate, cirq.WaitGate)
+
+    for op in circuit[4].operations:
         assert isinstance(op.gate, AnalogDetuneQubit)
diff --git a/cirq-google/cirq_google/study/symbol_util.py b/cirq-google/cirq_google/study/symbol_util.py
index 9d13d461120..dec4b35d07c 100644
--- a/cirq-google/cirq_google/study/symbol_util.py
+++ b/cirq-google/cirq_google/study/symbol_util.py
@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 from typing import AbstractSet, Any, TypeAlias
 
 import sympy
@@ -84,7 +86,7 @@ def _parameter_names_(self) -> AbstractSet[str]:
 
     def _resolve_parameters_(
         self, resolver: cirq.ParamResolverOrSimilarType, recursive: bool
-    ) -> float:
+    ) -> TunitForDurationNanos:
         resolver_ = cirq.ParamResolver(resolver)
         return TunitForDurationNanos(
             duration_nanos=direct_symbol_replacement(self.duration_nanos, resolver_)[tu.ns]

From 227d8dc3dddda4d14c4d1e42555a35351e9b4ec4 Mon Sep 17 00:00:00 2001
From: Bicheng Ying <bichengying@gmail.com>
Date: Tue, 15 Jul 2025 19:05:46 +0000
Subject: [PATCH 04/10] type check

---
 .../analog_experiments/analog_trajectory_util_test.py        | 5 ++++-
 .../analog_experiments/generic_analog_circuit_test.py        | 2 +-
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/cirq-google/cirq_google/experimental/analog_experiments/analog_trajectory_util_test.py b/cirq-google/cirq_google/experimental/analog_experiments/analog_trajectory_util_test.py
index eb0c01ef775..9fc07484d07 100644
--- a/cirq-google/cirq_google/experimental/analog_experiments/analog_trajectory_util_test.py
+++ b/cirq-google/cirq_google/experimental/analog_experiments/analog_trajectory_util_test.py
@@ -18,6 +18,7 @@
 
 import cirq
 from cirq_google.experimental.analog_experiments import analog_trajectory_util as atu
+from cirq_google.study import symbol_util as su
 
 
 @pytest.fixture
@@ -51,7 +52,9 @@ def test_duration_nanos_in_freq_map() -> None:
     fm = atu.FrequencyMap(10 * tu.ns, {}, {}, False)
     assert fm.duration_nanos() == 10
     fm = atu.FrequencyMap(sympy.Symbol("t"), {}, {}, False)
-    assert cirq.resolve_parameters(fm.duration_nanos(), {"t": 10 * tu.ns}).total_nanos() == 10
+    resolved_duration_nanos = cirq.resolve_parameters(fm.duration_nanos(), {"t": 10 * tu.ns})
+    assert isinstance(resolved_duration_nanos, su.TunitForDurationNanos)
+    assert resolved_duration_nanos.total_nanos() == 10
 
     with pytest.raises(ValueError, match="either be a tu.Value or a sympy.Symbol"):
         atu.FrequencyMap(10, {}, {}, False).duration_nanos()
diff --git a/cirq-google/cirq_google/experimental/analog_experiments/generic_analog_circuit_test.py b/cirq-google/cirq_google/experimental/analog_experiments/generic_analog_circuit_test.py
index bd1d8ddf5b9..3996c504a61 100644
--- a/cirq-google/cirq_google/experimental/analog_experiments/generic_analog_circuit_test.py
+++ b/cirq-google/cirq_google/experimental/analog_experiments/generic_analog_circuit_test.py
@@ -68,7 +68,7 @@ def test_make_one_moment_of_generic_analog_circuit() -> None:
     )
 
     trajectory = None  # we don't need trajector in this test.
-    builder = gac.GenericAnalogCircuitBuilder(trajectory)
+    builder = gac.GenericAnalogCircuitBuilder(trajectory)  # type: ignore
     moment = builder.make_one_moment(freq_map, prev_freq_map)
 
     assert len(moment.operations) == 5

From 8b53361f2fcb0fa50baa972f553bcbdd6d15f376 Mon Sep 17 00:00:00 2001
From: Bicheng Ying <bichengying@gmail.com>
Date: Wed, 23 Jul 2025 23:46:39 +0000
Subject: [PATCH 05/10] Remove the hacky part and add a new WaitGateWithUnits
 instead

---
 cirq-google/cirq_google/__init__.py           |  1 +
 .../analog_trajectory_util.py                 | 14 ----
 .../analog_trajectory_util_test.py            | 13 ----
 .../generic_analog_circuit.py                 | 22 ++-----
 .../cirq_google/json_resolver_cache.py        |  1 +
 .../json_test_data/WaitGateWithUnits.json     |  7 ++
 .../json_test_data/WaitGateWithUnits.repr     |  1 +
 cirq-google/cirq_google/ops/__init__.py       |  2 +
 cirq-google/cirq_google/ops/wait_gate.py      | 65 +++++++++++++++++++
 cirq-google/cirq_google/ops/wait_gate_test.py | 62 ++++++++++++++++++
 cirq-google/cirq_google/study/symbol_util.py  | 38 -----------
 .../cirq_google/study/symbol_util_test.py     | 18 -----
 12 files changed, 145 insertions(+), 99 deletions(-)
 create mode 100644 cirq-google/cirq_google/json_test_data/WaitGateWithUnits.json
 create mode 100644 cirq-google/cirq_google/json_test_data/WaitGateWithUnits.repr
 create mode 100644 cirq-google/cirq_google/ops/wait_gate.py
 create mode 100644 cirq-google/cirq_google/ops/wait_gate_test.py

diff --git a/cirq-google/cirq_google/__init__.py b/cirq-google/cirq_google/__init__.py
index 91250d859d3..81531df86e9 100644
--- a/cirq-google/cirq_google/__init__.py
+++ b/cirq-google/cirq_google/__init__.py
@@ -64,6 +64,7 @@
     PhysicalZTag as PhysicalZTag,
     SYC as SYC,
     SycamoreGate as SycamoreGate,
+    WaitGateWithUnit as WaitGateWithUnit,
     WILLOW as WILLOW,
     WillowGate as WillowGate,
 )
diff --git a/cirq-google/cirq_google/experimental/analog_experiments/analog_trajectory_util.py b/cirq-google/cirq_google/experimental/analog_experiments/analog_trajectory_util.py
index 8b617be65c3..9f60e034d74 100644
--- a/cirq-google/cirq_google/experimental/analog_experiments/analog_trajectory_util.py
+++ b/cirq-google/cirq_google/experimental/analog_experiments/analog_trajectory_util.py
@@ -19,7 +19,6 @@
 import attrs
 import matplotlib.pyplot as plt
 import numpy as np
-import sympy
 import tunits as tu
 
 import cirq
@@ -74,19 +73,6 @@ def _resolve_parameters_(
             is_wait_step=self.is_wait_step,
         )
 
-    def duration_nanos(self) -> float | su.TunitForDurationNanos:
-        # The following is the patching code for symbol/tunit can work
-        # with cirq.Duration object.
-        if isinstance(self.duration, tu.Value):
-            nanos = self.duration[tu.ns]
-        elif isinstance(self.duration, sympy.Symbol):
-            nanos = su.TunitForDurationNanos(self.duration)
-        else:
-            raise ValueError(
-                "The duration in the freq map must either be a tu.Value or a sympy.Symbol."
-            )
-        return nanos
-
 
 class AnalogTrajectory:
     """Class for handling qubit frequency and coupling trajectories that
diff --git a/cirq-google/cirq_google/experimental/analog_experiments/analog_trajectory_util_test.py b/cirq-google/cirq_google/experimental/analog_experiments/analog_trajectory_util_test.py
index 9fc07484d07..d55afd19de0 100644
--- a/cirq-google/cirq_google/experimental/analog_experiments/analog_trajectory_util_test.py
+++ b/cirq-google/cirq_google/experimental/analog_experiments/analog_trajectory_util_test.py
@@ -18,7 +18,6 @@
 
 import cirq
 from cirq_google.experimental.analog_experiments import analog_trajectory_util as atu
-from cirq_google.study import symbol_util as su
 
 
 @pytest.fixture
@@ -48,18 +47,6 @@ def test_freq_map_resolve(freq_map: atu.FrequencyMap) -> None:
     )
 
 
-def test_duration_nanos_in_freq_map() -> None:
-    fm = atu.FrequencyMap(10 * tu.ns, {}, {}, False)
-    assert fm.duration_nanos() == 10
-    fm = atu.FrequencyMap(sympy.Symbol("t"), {}, {}, False)
-    resolved_duration_nanos = cirq.resolve_parameters(fm.duration_nanos(), {"t": 10 * tu.ns})
-    assert isinstance(resolved_duration_nanos, su.TunitForDurationNanos)
-    assert resolved_duration_nanos.total_nanos() == 10
-
-    with pytest.raises(ValueError, match="either be a tu.Value or a sympy.Symbol"):
-        atu.FrequencyMap(10, {}, {}, False).duration_nanos()
-
-
 FreqMapType = tuple[tu.Value, dict[str, tu.Value | None], dict[tuple[str, str], tu.Value]]
 
 
diff --git a/cirq-google/cirq_google/experimental/analog_experiments/generic_analog_circuit.py b/cirq-google/cirq_google/experimental/analog_experiments/generic_analog_circuit.py
index a3421b0416e..1859d8cb788 100644
--- a/cirq-google/cirq_google/experimental/analog_experiments/generic_analog_circuit.py
+++ b/cirq-google/cirq_google/experimental/analog_experiments/generic_analog_circuit.py
@@ -17,7 +17,7 @@
 
 import cirq
 from cirq_google.experimental.analog_experiments import analog_trajectory_util as atu
-from cirq_google.ops import analog_detune_gates as adg
+from cirq_google.ops import analog_detune_gates as adg, wait_gate as wg
 from cirq_google.study import symbol_util as su
 
 
@@ -90,21 +90,11 @@ def make_circuit(self) -> cirq.Circuit:
         moments = []
         for freq_map in self.trajectory.full_trajectory[1:]:
             if freq_map.is_wait_step:
-                target = [_to_grid_qubit(q) for q in self.trajectory.qubits]
-                d = freq_map.duration_nanos()
-                if isinstance(d, float):
-                    wait_gate = cirq.WaitGate(
-                        cirq.Duration(nanos=d), qid_shape=cirq.qid_shape(target)
-                    )
-                else:
-                    # The following is patching solution for resolving the parameter
-                    # can be tunits. It should only work for pyle internal translation.
-                    wait_gate = cirq.WaitGate(
-                        cirq.Duration(nanos=1), qid_shape=cirq.qid_shape(target)
-                    )
-                    wait_gate._duration = d  # type: ignore
-
-                moment = cirq.Moment(wait_gate.on(*target))
+                targets = [_to_grid_qubit(q) for q in self.trajectory.qubits]
+                wait_gate = wg.WaitGateWithUnit(
+                    freq_map.duration, qid_shape=cirq.qid_shape(targets)
+                )
+                moment = cirq.Moment(wait_gate.on(*targets))
             else:
                 moment = self.make_one_moment(freq_map, prev_freq_map)
             moments.append(moment)
diff --git a/cirq-google/cirq_google/json_resolver_cache.py b/cirq-google/cirq_google/json_resolver_cache.py
index 0f2339c66bf..8df67ffb5f8 100644
--- a/cirq-google/cirq_google/json_resolver_cache.py
+++ b/cirq-google/cirq_google/json_resolver_cache.py
@@ -56,6 +56,7 @@ def _old_xmon(*args, **kwargs):
             cirq_google.experimental.PerQubitDepolarizingWithDampedReadoutNoiseModel
         ),
         'SycamoreGate': cirq_google.SycamoreGate,
+        'WaitGateWithUnit': cirq_google.WaitGateWithUnit,
         'WillowGate': cirq_google.WillowGate,
         # cirq_google.GateTabulation has been removed and replaced by cirq.TwoQubitGateTabulation.
         'GateTabulation': TwoQubitGateTabulation,
diff --git a/cirq-google/cirq_google/json_test_data/WaitGateWithUnits.json b/cirq-google/cirq_google/json_test_data/WaitGateWithUnits.json
new file mode 100644
index 00000000000..700e06e86a8
--- /dev/null
+++ b/cirq-google/cirq_google/json_test_data/WaitGateWithUnits.json
@@ -0,0 +1,7 @@
+{
+    "cirq_type": "WaitGateWithUnit",
+    "duration": {
+        "cirq_type": "sympy.Symbol",
+        "name": "d"
+    }
+}
\ No newline at end of file
diff --git a/cirq-google/cirq_google/json_test_data/WaitGateWithUnits.repr b/cirq-google/cirq_google/json_test_data/WaitGateWithUnits.repr
new file mode 100644
index 00000000000..e083ed001fe
--- /dev/null
+++ b/cirq-google/cirq_google/json_test_data/WaitGateWithUnits.repr
@@ -0,0 +1 @@
+cirq_google.WaitGateWithUnit(duration=sympy.Symbol("d"))
\ No newline at end of file
diff --git a/cirq-google/cirq_google/ops/__init__.py b/cirq-google/cirq_google/ops/__init__.py
index 13a70bbe5f8..565fcdcef5c 100644
--- a/cirq-google/cirq_google/ops/__init__.py
+++ b/cirq-google/cirq_google/ops/__init__.py
@@ -38,4 +38,6 @@
     DynamicalDecouplingTag as DynamicalDecouplingTag,
 )
 
+from cirq_google.ops.wait_gate import WaitGateWithUnit as WaitGateWithUnit
+
 from cirq_google.ops.willow_gate import WillowGate as WillowGate, WILLOW as WILLOW
diff --git a/cirq-google/cirq_google/ops/wait_gate.py b/cirq-google/cirq_google/ops/wait_gate.py
new file mode 100644
index 00000000000..84d5105483d
--- /dev/null
+++ b/cirq-google/cirq_google/ops/wait_gate.py
@@ -0,0 +1,65 @@
+# Copyright 2025 The Cirq Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 __future__ import annotations
+
+import sympy
+import tunits as tu
+
+import cirq
+from cirq_google.study import symbol_util as su
+
+
+class WaitGateWithUnit(cirq.WaitGate):
+    """A wrapper on top of WaitGate that can work with units."""
+
+    def __init__(
+        self,
+        duration: su.ValueOrSymbol,
+        num_qubits: int | None = None,
+        qid_shape: tuple[int, ...] | None = None,
+    ):
+        if not isinstance(duration, su.ValueOrSymbol):
+            raise ValueError("The duration must either be a tu.Value or a sympy.Symbol.")
+        self._duration = duration
+        # The rest is copy-pasted from WaitGate. We just cannot use
+        # super().__init__ because of the duration.
+        if qid_shape is None:
+            if num_qubits is None:
+                # Assume one qubit for backwards compatibility
+                qid_shape = (2,)
+            else:
+                qid_shape = (2,) * num_qubits
+        if num_qubits is None:
+            num_qubits = len(qid_shape)
+        if not qid_shape:
+            raise ValueError('Waiting on an empty set of qubits.')
+        if num_qubits != len(qid_shape):
+            raise ValueError('len(qid_shape) != num_qubits')
+        self._qid_shape = qid_shape
+
+    @property
+    def duration(self) -> sympy.Symbol | cirq.Duration:
+        if isinstance(self._duration, sympy.Symbol):
+            return self._duration
+        return cirq.Duration(nanos=self._duration[tu.ns])
+
+    def _resolve_parameters_(
+        self, resolver: cirq.ParamResolver, recursive: bool
+    ) -> WaitGateWithUnit:
+        if isinstance(self._duration, sympy.Symbol):
+            resolver_ = cirq.ParamResolver(resolver)
+            _duration = su.direct_symbol_replacement(self._duration, resolver_)
+            return WaitGateWithUnit(_duration, qid_shape=self._qid_shape)
+        return self
diff --git a/cirq-google/cirq_google/ops/wait_gate_test.py b/cirq-google/cirq_google/ops/wait_gate_test.py
new file mode 100644
index 00000000000..829fc8a3b2c
--- /dev/null
+++ b/cirq-google/cirq_google/ops/wait_gate_test.py
@@ -0,0 +1,62 @@
+# Copyright 2025 The Cirq Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+import pytest
+import sympy
+import tunits as tu
+
+import cirq
+from cirq_google.ops import wait_gate as wg
+
+
+def test_wait_gate_with_unit_init() -> None:
+    g = wg.WaitGateWithUnit(1 * tu.us)
+    assert g.duration == cirq.Duration(nanos=1000)
+
+    g = wg.WaitGateWithUnit(1 * tu.us, num_qubits=2)
+    assert g._qid_shape == (2, 2)
+
+    g = wg.WaitGateWithUnit(sympy.Symbol("d"))
+    assert g.duration == sympy.Symbol("d")
+
+    with pytest.raises(ValueError, match="either be a tu.Value or a sympy.Symbol."):
+        wg.WaitGateWithUnit(10)
+
+    with pytest.raises(ValueError, match="Waiting on an empty set of qubits."):
+        wg.WaitGateWithUnit(10 * tu.ns, qid_shape=())
+
+    with pytest.raises(ValueError, match="num_qubits"):
+        wg.WaitGateWithUnit(10 * tu.ns, qid_shape=(2, 2), num_qubits=5)
+
+
+def test_wait_gate_with_units_resolving() -> None:
+    gate = wg.WaitGateWithUnit(sympy.Symbol("d"))
+
+    resolved_gate = cirq.resolve_parameters(gate, {"d": 10 * tu.ns})
+    assert resolved_gate.duration == cirq.Duration(nanos=10)
+
+
+def test_wait_gate_equality() -> None:
+    gate1 = wg.WaitGateWithUnit(10 * tu.ns)
+    gate2 = wg.WaitGateWithUnit(10 * tu.ns)
+    assert gate1 == gate2
+
+    gate_symbol_1 = wg.WaitGateWithUnit(sympy.Symbol("a"))
+    gate_symbol_2 = wg.WaitGateWithUnit(sympy.Symbol("a"))
+    assert gate_symbol_1 == gate_symbol_2
+    assert gate_symbol_1 != gate1
+
+
+def test_wait_gate_jsonify() -> None:
+    gate = wg.WaitGateWithUnit(sympy.Symbol("d"))
+    assert gate == cirq.read_json(json_text=cirq.to_json(gate))
diff --git a/cirq-google/cirq_google/study/symbol_util.py b/cirq-google/cirq_google/study/symbol_util.py
index dec4b35d07c..75b67ba2e29 100644
--- a/cirq-google/cirq_google/study/symbol_util.py
+++ b/cirq-google/cirq_google/study/symbol_util.py
@@ -53,41 +53,3 @@ def is_parameterized_dict(dict_with_value: dict[Any, ValueOrSymbol] | None) -> b
     if dict_with_value is None:
         return False  # pragma: no cover
     return any(cirq.is_parameterized(v) for v in dict_with_value.values())
-
-
-class TunitForDurationNanos:
-    """A wrapper class that can be used with symbols for duration nanos.
-
-    When resolving it, it will nanos (as float) so that we can use it
-    as the input of `cirq.Duration(nanos=TunitForDurationNanos())`.
-    """
-
-    def __init__(self, duration_nanos: sympy.Symbol | float):
-        self.duration_nanos = duration_nanos
-
-    def total_micros(self):
-        if isinstance(self.duration_nanos, sympy.Symbol):
-            return self.duration_nanos
-        return self.duration_nanos * 1000
-
-    def total_nanos(self):
-        return self.duration_nanos
-
-    def __repr__(self) -> str:
-        if isinstance(self.duration_nanos, sympy.Symbol):
-            return f"duration={self.duration_nanos}"
-        return f"duration={self.duration_nanos} ns"
-
-    def _is_parameterized_(self) -> bool:
-        return cirq.is_parameterized(self.duration_nanos)
-
-    def _parameter_names_(self) -> AbstractSet[str]:
-        return cirq.parameter_names(self.duration_nanos)
-
-    def _resolve_parameters_(
-        self, resolver: cirq.ParamResolverOrSimilarType, recursive: bool
-    ) -> TunitForDurationNanos:
-        resolver_ = cirq.ParamResolver(resolver)
-        return TunitForDurationNanos(
-            duration_nanos=direct_symbol_replacement(self.duration_nanos, resolver_)[tu.ns]
-        )
diff --git a/cirq-google/cirq_google/study/symbol_util_test.py b/cirq-google/cirq_google/study/symbol_util_test.py
index 9048237dde1..a6ef9d50b4d 100644
--- a/cirq-google/cirq_google/study/symbol_util_test.py
+++ b/cirq-google/cirq_google/study/symbol_util_test.py
@@ -46,21 +46,3 @@ def test_direct_symbol_replacement():
     value_resolved = [su.direct_symbol_replacement(v, resolver) for v in value_list]
 
     assert value_resolved == [5 * tu.ns, 8 * tu.GHz, sympy.Symbol("z"), 123, "fd"]
-
-
-def test_tunit_for_duration_nanos():
-    d = su.TunitForDurationNanos(duration_nanos=sympy.Symbol("t"))
-    assert str(d) == "duration=t"
-    assert cirq.is_parameterized(d)
-    assert d._parameter_names_() == {"t"}
-    assert cirq.resolve_parameters(d, {"t": 5 * tu.ns}).duration_nanos == 5.0
-
-
-def test_tunit_for_duration_total_nanos():
-    d = su.TunitForDurationNanos(duration_nanos=sympy.Symbol("t"))
-    assert d.total_nanos() == sympy.Symbol("t")
-    assert d.total_micros() == sympy.Symbol("t")
-    d_resolved = cirq.resolve_parameters(d, {"t": 5 * tu.ns})
-    assert str(d_resolved) == "duration=5.0 ns"
-    assert d_resolved.total_nanos() == 5
-    assert d_resolved.total_micros() == 5000

From 317d5125d24cb0718f245462f1f78ddbaae701cd Mon Sep 17 00:00:00 2001
From: Bicheng Ying <bichengying@gmail.com>
Date: Wed, 23 Jul 2025 23:53:19 +0000
Subject: [PATCH 06/10] rename

---
 .../{WaitGateWithUnits.json => WaitGateWithUnit.json}          | 0
 .../{WaitGateWithUnits.repr => WaitGateWithUnit.repr}          | 0
 cirq-google/cirq_google/ops/wait_gate.py                       | 3 ++-
 3 files changed, 2 insertions(+), 1 deletion(-)
 rename cirq-google/cirq_google/json_test_data/{WaitGateWithUnits.json => WaitGateWithUnit.json} (100%)
 rename cirq-google/cirq_google/json_test_data/{WaitGateWithUnits.repr => WaitGateWithUnit.repr} (100%)

diff --git a/cirq-google/cirq_google/json_test_data/WaitGateWithUnits.json b/cirq-google/cirq_google/json_test_data/WaitGateWithUnit.json
similarity index 100%
rename from cirq-google/cirq_google/json_test_data/WaitGateWithUnits.json
rename to cirq-google/cirq_google/json_test_data/WaitGateWithUnit.json
diff --git a/cirq-google/cirq_google/json_test_data/WaitGateWithUnits.repr b/cirq-google/cirq_google/json_test_data/WaitGateWithUnit.repr
similarity index 100%
rename from cirq-google/cirq_google/json_test_data/WaitGateWithUnits.repr
rename to cirq-google/cirq_google/json_test_data/WaitGateWithUnit.repr
diff --git a/cirq-google/cirq_google/ops/wait_gate.py b/cirq-google/cirq_google/ops/wait_gate.py
index 84d5105483d..071fcbaabd2 100644
--- a/cirq-google/cirq_google/ops/wait_gate.py
+++ b/cirq-google/cirq_google/ops/wait_gate.py
@@ -32,7 +32,8 @@ def __init__(
     ):
         if not isinstance(duration, su.ValueOrSymbol):
             raise ValueError("The duration must either be a tu.Value or a sympy.Symbol.")
-        self._duration = duration
+        # Override the original duration
+        self._duration: su.ValueOrSymbol = duration
         # The rest is copy-pasted from WaitGate. We just cannot use
         # super().__init__ because of the duration.
         if qid_shape is None:

From 6522275f5f4b3079e4630b3e4e416ef035649c64 Mon Sep 17 00:00:00 2001
From: Bicheng Ying <bichengying@gmail.com>
Date: Wed, 23 Jul 2025 23:57:45 +0000
Subject: [PATCH 07/10] type

---
 cirq-google/cirq_google/ops/wait_gate.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/cirq-google/cirq_google/ops/wait_gate.py b/cirq-google/cirq_google/ops/wait_gate.py
index 071fcbaabd2..ef022946bb6 100644
--- a/cirq-google/cirq_google/ops/wait_gate.py
+++ b/cirq-google/cirq_google/ops/wait_gate.py
@@ -33,7 +33,8 @@ def __init__(
         if not isinstance(duration, su.ValueOrSymbol):
             raise ValueError("The duration must either be a tu.Value or a sympy.Symbol.")
         # Override the original duration
-        self._duration: su.ValueOrSymbol = duration
+        self._duration: su.ValueOrSymbol = duration  # type: ignore[assignment]
+
         # The rest is copy-pasted from WaitGate. We just cannot use
         # super().__init__ because of the duration.
         if qid_shape is None:

From 61ce43abec0e83ac064d79eecd4f646b5c094020 Mon Sep 17 00:00:00 2001
From: Bicheng Ying <bichengying@gmail.com>
Date: Thu, 24 Jul 2025 19:49:08 +0000
Subject: [PATCH 08/10] coverage

---
 cirq-google/cirq_google/ops/wait_gate_test.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/cirq-google/cirq_google/ops/wait_gate_test.py b/cirq-google/cirq_google/ops/wait_gate_test.py
index 829fc8a3b2c..3d23694f8c5 100644
--- a/cirq-google/cirq_google/ops/wait_gate_test.py
+++ b/cirq-google/cirq_google/ops/wait_gate_test.py
@@ -44,6 +44,10 @@ def test_wait_gate_with_units_resolving() -> None:
 
     resolved_gate = cirq.resolve_parameters(gate, {"d": 10 * tu.ns})
     assert resolved_gate.duration == cirq.Duration(nanos=10)
+    
+    gate = wg.WaitGateWithUnit(10*tu.ns)
+    assert gate._resolve_parameters_({}, True) == gate
+
 
 
 def test_wait_gate_equality() -> None:

From f6be921bd0c3a1ea4a29c89daaee680223671db9 Mon Sep 17 00:00:00 2001
From: Bicheng Ying <bichengying@gmail.com>
Date: Thu, 24 Jul 2025 19:50:50 +0000
Subject: [PATCH 09/10] format

---
 cirq-google/cirq_google/ops/wait_gate_test.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/cirq-google/cirq_google/ops/wait_gate_test.py b/cirq-google/cirq_google/ops/wait_gate_test.py
index 3d23694f8c5..d12c7acf7e1 100644
--- a/cirq-google/cirq_google/ops/wait_gate_test.py
+++ b/cirq-google/cirq_google/ops/wait_gate_test.py
@@ -44,10 +44,9 @@ def test_wait_gate_with_units_resolving() -> None:
 
     resolved_gate = cirq.resolve_parameters(gate, {"d": 10 * tu.ns})
     assert resolved_gate.duration == cirq.Duration(nanos=10)
-    
-    gate = wg.WaitGateWithUnit(10*tu.ns)
-    assert gate._resolve_parameters_({}, True) == gate
 
+    gate = wg.WaitGateWithUnit(10 * tu.ns)
+    assert gate._resolve_parameters_({}, True) == gate
 
 
 def test_wait_gate_equality() -> None:

From 5b096619a9baa9e559fc4154c09f7c0df3d4f1f6 Mon Sep 17 00:00:00 2001
From: Bicheng Ying <bichengying@gmail.com>
Date: Fri, 25 Jul 2025 00:26:12 +0000
Subject: [PATCH 10/10] typecheck

---
 cirq-google/cirq_google/ops/wait_gate.py      | 3 +--
 cirq-google/cirq_google/ops/wait_gate_test.py | 2 +-
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/cirq-google/cirq_google/ops/wait_gate.py b/cirq-google/cirq_google/ops/wait_gate.py
index ef022946bb6..72aa1cae750 100644
--- a/cirq-google/cirq_google/ops/wait_gate.py
+++ b/cirq-google/cirq_google/ops/wait_gate.py
@@ -61,7 +61,6 @@ def _resolve_parameters_(
         self, resolver: cirq.ParamResolver, recursive: bool
     ) -> WaitGateWithUnit:
         if isinstance(self._duration, sympy.Symbol):
-            resolver_ = cirq.ParamResolver(resolver)
-            _duration = su.direct_symbol_replacement(self._duration, resolver_)
+            _duration = su.direct_symbol_replacement(self._duration, resolver)
             return WaitGateWithUnit(_duration, qid_shape=self._qid_shape)
         return self
diff --git a/cirq-google/cirq_google/ops/wait_gate_test.py b/cirq-google/cirq_google/ops/wait_gate_test.py
index d12c7acf7e1..f8c0690b55f 100644
--- a/cirq-google/cirq_google/ops/wait_gate_test.py
+++ b/cirq-google/cirq_google/ops/wait_gate_test.py
@@ -46,7 +46,7 @@ def test_wait_gate_with_units_resolving() -> None:
     assert resolved_gate.duration == cirq.Duration(nanos=10)
 
     gate = wg.WaitGateWithUnit(10 * tu.ns)
-    assert gate._resolve_parameters_({}, True) == gate
+    assert gate._resolve_parameters_(cirq.ParamResolver({}), True) == gate
 
 
 def test_wait_gate_equality() -> None: