cirq-core 1.5.0.dev20250409222543__py3-none-any.whl → 1.6.0__py3-none-any.whl

cirq/work/observable_readout_calibration_test.py

@@ -1,4 +1,7 @@
 # pylint: disable=wrong-or-nonexistent-copyright-notice
+
+from __future__ import annotations
+
 from typing import Sequence
 
 import numpy as np
@@ -19,7 +22,7 @@ class DepolarizingWithDampedReadoutNoiseModel(cirq.NoiseModel):
         self.readout_noise_gate = cirq.BitFlipChannel(bitflip_prob)
         self.readout_decay_gate = cirq.AmplitudeDampingChannel(decay_prob)
 
-    def noisy_moment(self, moment: 'cirq.Moment', system_qubits: Sequence['cirq.Qid']):
+    def noisy_moment(self, moment: cirq.Moment, system_qubits: Sequence[cirq.Qid]):
         if cirq.devices.noise_model.validate_all_measurements(moment):
             return [
                 cirq.Moment(self.readout_decay_gate(q) for q in system_qubits),
@@ -30,7 +33,7 @@ class DepolarizingWithDampedReadoutNoiseModel(cirq.NoiseModel):
             return [moment, cirq.Moment(self.qubit_noise_gate(q) for q in system_qubits)]
 
 
-def test_calibrate_readout_error():
+def test_calibrate_readout_error() -> None:
     sampler = cirq.DensityMatrixSimulator(
         noise=DepolarizingWithDampedReadoutNoiseModel(
             depol_prob=1e-3, bitflip_prob=0.03, decay_prob=0.03

cirq/work/observable_settings.py

@@ -12,19 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import dataclasses
 import numbers
-from typing import (
-    AbstractSet,
-    Dict,
-    FrozenSet,
-    Iterable,
-    Mapping,
-    Optional,
-    Tuple,
-    TYPE_CHECKING,
-    Union,
-)
+from typing import Iterable, Mapping, TYPE_CHECKING
 
 import sympy
 
@@ -72,7 +64,7 @@ class InitObsSetting:
         return protocols.dataclass_json_dict(self)
 
 
-def _max_weight_observable(observables: Iterable[ops.PauliString]) -> Optional[ops.PauliString]:
+def _max_weight_observable(observables: Iterable[ops.PauliString]) -> ops.PauliString | None:
     """Create a new observable that is compatible with all input observables
     and has the maximum non-identity elements.
 
@@ -88,7 +80,7 @@ def _max_weight_observable(observables: Iterable[ops.PauliString]) -> Optional[o
     The returned value need not actually be present in the input observables.
     Coefficients from input observables will be dropped.
     """
-    qubit_pauli_map: Dict[ops.Qid, ops.Pauli] = {}
+    qubit_pauli_map: dict[ops.Qid, ops.Pauli] = {}
     for observable in observables:
         for qubit, pauli in observable.items():
             if qubit in qubit_pauli_map:
@@ -99,7 +91,7 @@ def _max_weight_observable(observables: Iterable[ops.PauliString]) -> Optional[o
     return ops.PauliString(qubit_pauli_map)
 
 
-def _max_weight_state(states: Iterable[value.ProductState]) -> Optional[value.ProductState]:
+def _max_weight_state(states: Iterable[value.ProductState]) -> value.ProductState | None:
     """Create a new state that is compatible with all input states
     and has the maximum weight.
 
@@ -112,7 +104,7 @@ def _max_weight_state(states: Iterable[value.ProductState]) -> Optional[value.Pr
     "+X(0) * -Z(1)". Asking for the max weight state of something like
     [+X(0), +Z(0)] will return None.
     """
-    qubit_state_map: Dict[ops.Qid, _NamedOneQubitState] = {}
+    qubit_state_map: dict[ops.Qid, _NamedOneQubitState] = {}
     for state in states:
         for qubit, named_state in state:
             if qubit in qubit_state_map:
@@ -123,13 +115,13 @@ def _max_weight_state(states: Iterable[value.ProductState]) -> Optional[value.Pr
     return value.ProductState(qubit_state_map)
 
 
-def zeros_state(qubits: Iterable['cirq.Qid']):
+def zeros_state(qubits: Iterable[cirq.Qid]):
     """Return the ProductState that is |00..00> on all qubits."""
     return value.ProductState({q: value.KET_ZERO for q in qubits})
 
 
 def observables_to_settings(
-    observables: Iterable['cirq.PauliString'], qubits: Iterable['cirq.Qid']
+    observables: Iterable[cirq.PauliString], qubits: Iterable[cirq.Qid]
 ) -> Iterable[InitObsSetting]:
     """Transform an observable to an InitObsSetting initialized in the
     all-zeros state.
@@ -138,7 +130,7 @@ def observables_to_settings(
         yield InitObsSetting(init_state=zeros_state(qubits), observable=observable)
 
 
-def _fix_precision(val: Union[value.Scalar, sympy.Expr], precision) -> Union[int, Tuple[int, int]]:
+def _fix_precision(val: value.Scalar | sympy.Expr, precision) -> int | tuple[int, int]:
     """Convert floating point or complex numbers to (implicitly) fixed point
     integers. Complex numbers will return fixed-point (real, imag) tuples.
 
@@ -154,9 +146,8 @@ def _fix_precision(val: Union[value.Scalar, sympy.Expr], precision) -> Union[int
 
 
 def _hashable_param(
-    param_tuples: AbstractSet[Tuple[Union[str, sympy.Expr], Union[value.Scalar, sympy.Expr]]],
-    precision=1e7,
-) -> FrozenSet[Tuple[str, Union[int, Tuple[int, int]]]]:
+    param_tuples: Iterable[tuple[str | sympy.Expr, value.Scalar | sympy.Expr]], precision=1e7
+) -> frozenset[tuple[str, int | tuple[int, int]]]:
     """Hash circuit parameters using fixed precision.
 
     Circuit parameters can be complex but we also need to use them as
@@ -178,7 +169,7 @@ class _MeasurementSpec:
     """
 
     max_setting: InitObsSetting
-    circuit_params: Mapping[Union[str, sympy.Expr], Union[value.Scalar, sympy.Expr]]
+    circuit_params: Mapping[str | sympy.Expr, value.Scalar | sympy.Expr]
 
     def __hash__(self):
         return hash((self.max_setting, _hashable_param(self.circuit_params.items())))

cirq/work/observable_settings_test.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import pytest
 import sympy
 
@@ -20,7 +22,7 @@ from cirq.work import _MeasurementSpec, InitObsSetting, observables_to_settings
 from cirq.work.observable_settings import _hashable_param, _max_weight_observable, _max_weight_state
 
 
-def test_init_obs_setting():
+def test_init_obs_setting() -> None:
     q0, q1 = cirq.LineQubit.range(2)
     setting = InitObsSetting(
         init_state=cirq.KET_ZERO(q0) * cirq.KET_ZERO(q1), observable=cirq.X(q0) * cirq.Y(q1)
@@ -32,7 +34,7 @@ def test_init_obs_setting():
         setting = InitObsSetting(init_state=cirq.KET_ZERO(q0), observable=cirq.X(q0) * cirq.Y(q1))
 
 
-def test_max_weight_observable():
+def test_max_weight_observable() -> None:
     q0, q1 = cirq.LineQubit.range(2)
     observables = [cirq.X(q0), cirq.X(q1)]
     assert _max_weight_observable(observables) == cirq.X(q0) * cirq.X(q1)
@@ -41,7 +43,7 @@ def test_max_weight_observable():
     assert _max_weight_observable(observables) is None
 
 
-def test_max_weight_state():
+def test_max_weight_state() -> None:
     q0, q1 = cirq.LineQubit.range(2)
     states = [cirq.KET_PLUS(q0), cirq.KET_PLUS(q1)]
     assert _max_weight_state(states) == cirq.KET_PLUS(q0) * cirq.KET_PLUS(q1)
@@ -50,7 +52,7 @@ def test_max_weight_state():
     assert _max_weight_state(states) is None
 
 
-def test_observable_to_setting():
+def test_observable_to_setting() -> None:
     q0, q1, q2 = cirq.LineQubit.range(3)
     observables = [cirq.X(q0) * cirq.Y(q1), cirq.Z(q2) * 1]
 
@@ -62,7 +64,7 @@ def test_observable_to_setting():
     assert list(observables_to_settings(observables, qubits=[q0, q1, q2])) == settings_should_be
 
 
-def test_param_hash():
+def test_param_hash() -> None:
     params1 = [('beta', 1.23), ('gamma', 4.56)]
     params2 = [('beta', 1.23), ('gamma', 4.56)]
     params3 = [('beta', 1.24), ('gamma', 4.57)]
@@ -81,7 +83,7 @@ def test_param_hash():
     assert hash(_hashable_param(params4)) == hash(_hashable_param(params5))
 
 
-def test_measurement_spec():
+def test_measurement_spec() -> None:
     q0, q1 = cirq.LineQubit.range(2)
     setting = InitObsSetting(
         init_state=cirq.KET_ZERO(q0) * cirq.KET_ZERO(q1), observable=cirq.X(q0) * cirq.Y(q1)
@@ -96,7 +98,7 @@ def test_measurement_spec():
     cirq.testing.assert_equivalent_repr(meas_spec)
 
 
-def test_measurement_spec_no_symbols():
+def test_measurement_spec_no_symbols() -> None:
     q0, q1 = cirq.LineQubit.range(2)
     setting = InitObsSetting(
         init_state=cirq.KET_ZERO(q0) * cirq.KET_ZERO(q1), observable=cirq.X(q0) * cirq.Y(q1)

cirq/work/pauli_sum_collector.py

@@ -12,8 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import collections
-from typing import cast, Dict, Optional, TYPE_CHECKING, Union
+from typing import cast, TYPE_CHECKING
 
 import numpy as np
 
@@ -29,8 +31,8 @@ class PauliSumCollector(collector.Collector):
 
     def __init__(
         self,
-        circuit: 'cirq.AbstractCircuit',
-        observable: 'cirq.PauliSumLike',
+        circuit: cirq.AbstractCircuit,
+        observable: cirq.PauliSumLike,
         *,
         samples_per_term: int,
         max_samples_per_job: int = 1000000,
@@ -58,12 +60,12 @@ class PauliSumCollector(collector.Collector):
             if not p:
                 self._identity_offset += p.coefficient
 
-        self._zeros: Dict[ops.PauliString, int] = collections.defaultdict(lambda: 0)
-        self._ones: Dict[ops.PauliString, int] = collections.defaultdict(lambda: 0)
+        self._zeros: dict[ops.PauliString, int] = collections.defaultdict(lambda: 0)
+        self._ones: dict[ops.PauliString, int] = collections.defaultdict(lambda: 0)
         self._samples_per_term = samples_per_term
         self._total_samples_requested = 0
 
-    def next_job(self) -> Optional['cirq.CircuitSampleJob']:
+    def next_job(self) -> cirq.CircuitSampleJob | None:
         i = self._total_samples_requested // self._samples_per_term
         if i >= len(self._pauli_coef_terms):
             return None
@@ -77,13 +79,13 @@ class PauliSumCollector(collector.Collector):
             tag=pauli,
         )
 
-    def on_job_result(self, job: 'cirq.CircuitSampleJob', result: 'cirq.Result'):
+    def on_job_result(self, job: cirq.CircuitSampleJob, result: cirq.Result):
         job_id = cast(ops.PauliString, job.tag)
         parities = result.histogram(key='out', fold_func=lambda bits: np.sum(bits) % 2)
         self._zeros[job_id] += parities[0]
         self._ones[job_id] += parities[1]
 
-    def estimated_energy(self) -> Union[float, complex]:
+    def estimated_energy(self) -> float | complex:
         """Sums up the sampled expectations, weighted by their coefficients."""
         energy = 0j
         for pauli_string, coef in self._pauli_coef_terms:
@@ -99,8 +101,8 @@ class PauliSumCollector(collector.Collector):
 
 
 def _circuit_plus_pauli_string_measurements(
-    circuit: 'cirq.AbstractCircuit', pauli_string: 'cirq.PauliString'
-) -> 'cirq.AbstractCircuit':
+    circuit: cirq.AbstractCircuit, pauli_string: cirq.PauliString
+) -> cirq.AbstractCircuit:
     """A circuit measuring the given observable at the end of the given circuit."""
     assert pauli_string
     return circuit.from_moments(

cirq/work/pauli_sum_collector_test.py

@@ -12,13 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import duet
 
 import cirq
 
 
 @duet.sync
-async def test_pauli_string_sample_collector():
+async def test_pauli_string_sample_collector() -> None:
     a, b = cirq.LineQubit.range(2)
     p = cirq.PauliSumCollector(
         circuit=cirq.Circuit(cirq.H(a), cirq.CNOT(a, b), cirq.X(a), cirq.Z(b)),
@@ -28,26 +30,24 @@ async def test_pauli_string_sample_collector():
         + (1 - 0j),
         samples_per_term=100,
     )
-    result = await p.collect_async(sampler=cirq.Simulator())
-    assert result is None
+    await p.collect_async(sampler=cirq.Simulator())
     energy = p.estimated_energy()
     assert isinstance(energy, float) and energy == 12
 
 
 @duet.sync
-async def test_pauli_string_sample_single():
+async def test_pauli_string_sample_single() -> None:
     a, b = cirq.LineQubit.range(2)
     p = cirq.PauliSumCollector(
         circuit=cirq.Circuit(cirq.H(a), cirq.CNOT(a, b), cirq.X(a), cirq.Z(b)),
         observable=cirq.X(a) * cirq.X(b),
         samples_per_term=100,
     )
-    result = await p.collect_async(sampler=cirq.Simulator())
-    assert result is None
+    await p.collect_async(sampler=cirq.Simulator())
     assert p.estimated_energy() == -1
 
 
-def test_pauli_string_sample_collector_identity():
+def test_pauli_string_sample_collector_identity() -> None:
     p = cirq.PauliSumCollector(
         circuit=cirq.Circuit(), observable=cirq.PauliSum() + 2j, samples_per_term=100
     )
@@ -55,7 +55,7 @@ def test_pauli_string_sample_collector_identity():
     assert p.estimated_energy() == 2j
 
 
-def test_pauli_string_sample_collector_extra_qubit_z():
+def test_pauli_string_sample_collector_extra_qubit_z() -> None:
     a, b = cirq.LineQubit.range(2)
     p = cirq.PauliSumCollector(
         circuit=cirq.Circuit(cirq.H(a)), observable=3 * cirq.Z(b), samples_per_term=100
@@ -64,7 +64,7 @@ def test_pauli_string_sample_collector_extra_qubit_z():
     assert p.estimated_energy() == 3
 
 
-def test_pauli_string_sample_collector_extra_qubit_x():
+def test_pauli_string_sample_collector_extra_qubit_x() -> None:
     a, b = cirq.LineQubit.range(2)
     p = cirq.PauliSumCollector(
         circuit=cirq.Circuit(cirq.H(a)), observable=3 * cirq.X(b), samples_per_term=10000

cirq/work/sampler.py

@@ -11,10 +11,13 @@
 # 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.
+
 """Abstract base class for things sampling quantum circuits."""
 
+from __future__ import annotations
+
 import collections
-from typing import Dict, FrozenSet, List, Optional, Sequence, Tuple, TYPE_CHECKING, TypeVar, Union
+from typing import Sequence, TYPE_CHECKING, TypeVar
 
 import duet
 import pandas as pd
@@ -38,10 +41,10 @@ class Sampler(metaclass=value.ABCMetaImplementAnyOneOf):
 
     def run(
         self,
-        program: 'cirq.AbstractCircuit',
-        param_resolver: 'cirq.ParamResolverOrSimilarType' = None,
+        program: cirq.AbstractCircuit,
+        param_resolver: cirq.ParamResolverOrSimilarType = None,
         repetitions: int = 1,
-    ) -> 'cirq.Result':
+    ) -> cirq.Result:
         """Samples from the given `Circuit`.
 
         This mode of operation for a sampler will provide results
@@ -66,10 +69,10 @@ class Sampler(metaclass=value.ABCMetaImplementAnyOneOf):
 
     async def run_async(
         self,
-        program: 'cirq.AbstractCircuit',
-        param_resolver: 'cirq.ParamResolverOrSimilarType' = None,
+        program: cirq.AbstractCircuit,
+        param_resolver: cirq.ParamResolverOrSimilarType = None,
         repetitions: int = 1,
-    ) -> 'cirq.Result':
+    ) -> cirq.Result:
         """Asynchronously samples from the given Circuit.
 
         Provides measurement outcomes as a `cirq.Result` object.  This
@@ -88,12 +91,8 @@ class Sampler(metaclass=value.ABCMetaImplementAnyOneOf):
         return results[0]
 
     def sample(
-        self,
-        program: 'cirq.AbstractCircuit',
-        *,
-        repetitions: int = 1,
-        params: 'cirq.Sweepable' = None,
-    ) -> 'pd.DataFrame':
+        self, program: cirq.AbstractCircuit, *, repetitions: int = 1, params: cirq.Sweepable = None
+    ) -> pd.DataFrame:
         """Samples the given Circuit, producing a pandas data frame.
 
         This interface will operate in a similar way to the `run` method
@@ -180,21 +179,21 @@ class Sampler(metaclass=value.ABCMetaImplementAnyOneOf):
         return pd.concat(results)
 
     def _run_sweep_impl(
-        self, program: 'cirq.AbstractCircuit', params: 'cirq.Sweepable', repetitions: int = 1
-    ) -> Sequence['cirq.Result']:
+        self, program: cirq.AbstractCircuit, params: cirq.Sweepable, repetitions: int = 1
+    ) -> Sequence[cirq.Result]:
         """Implements run_sweep using run_sweep_async"""
         return duet.run(self.run_sweep_async, program, params, repetitions)
 
     async def _run_sweep_async_impl(
-        self, program: 'cirq.AbstractCircuit', params: 'cirq.Sweepable', repetitions: int = 1
-    ) -> Sequence['cirq.Result']:
+        self, program: cirq.AbstractCircuit, params: cirq.Sweepable, repetitions: int = 1
+    ) -> Sequence[cirq.Result]:
         """Implements run_sweep_async using run_sweep"""
         return self.run_sweep(program, params=params, repetitions=repetitions)
 
     @value.alternative(requires='run_sweep_async', implementation=_run_sweep_impl)
     def run_sweep(
-        self, program: 'cirq.AbstractCircuit', params: 'cirq.Sweepable', repetitions: int = 1
-    ) -> Sequence['cirq.Result']:
+        self, program: cirq.AbstractCircuit, params: cirq.Sweepable, repetitions: int = 1
+    ) -> Sequence[cirq.Result]:
         """Samples from the given Circuit.
 
         This allows for sweeping over different parameter values,
@@ -217,8 +216,8 @@ class Sampler(metaclass=value.ABCMetaImplementAnyOneOf):
 
     @value.alternative(requires='run_sweep', implementation=_run_sweep_async_impl)
     async def run_sweep_async(
-        self, program: 'cirq.AbstractCircuit', params: 'cirq.Sweepable', repetitions: int = 1
-    ) -> Sequence['cirq.Result']:
+        self, program: cirq.AbstractCircuit, params: cirq.Sweepable, repetitions: int = 1
+    ) -> Sequence[cirq.Result]:
         """Asynchronously samples from the given Circuit.
 
         By default, this method invokes `run_sweep` synchronously and simply
@@ -237,10 +236,10 @@ class Sampler(metaclass=value.ABCMetaImplementAnyOneOf):
 
     async def run_batch_async(
         self,
-        programs: Sequence['cirq.AbstractCircuit'],
-        params_list: Optional[Sequence['cirq.Sweepable']] = None,
-        repetitions: Union[int, Sequence[int]] = 1,
-    ) -> Sequence[Sequence['cirq.Result']]:
+        programs: Sequence[cirq.AbstractCircuit],
+        params_list: Sequence[cirq.Sweepable] | None = None,
+        repetitions: int | Sequence[int] = 1,
+    ) -> Sequence[Sequence[cirq.Result]]:
         """Runs the supplied circuits asynchronously.
 
         Each circuit provided in `programs` will pair with the optional
@@ -288,10 +287,10 @@ class Sampler(metaclass=value.ABCMetaImplementAnyOneOf):
 
     def _normalize_batch_args(
         self,
-        programs: Sequence['cirq.AbstractCircuit'],
-        params_list: Optional[Sequence['cirq.Sweepable']] = None,
-        repetitions: Union[int, Sequence[int]] = 1,
-    ) -> Tuple[Sequence['cirq.Sweepable'], Sequence[int]]:
+        programs: Sequence[cirq.AbstractCircuit],
+        params_list: Sequence[cirq.Sweepable] | None = None,
+        repetitions: int | Sequence[int] = 1,
+    ) -> tuple[Sequence[cirq.Sweepable], Sequence[int]]:
         if params_list is None:
             params_list = [None] * len(programs)
         if len(programs) != len(params_list):
@@ -310,11 +309,11 @@ class Sampler(metaclass=value.ABCMetaImplementAnyOneOf):
 
     def sample_expectation_values(
         self,
-        program: 'cirq.AbstractCircuit',
-        observables: Union['cirq.PauliSumLike', List['cirq.PauliSumLike']],
+        program: cirq.AbstractCircuit,
+        observables: cirq.PauliSumLike | list[cirq.PauliSumLike],
         *,
         num_samples: int,
-        params: 'cirq.Sweepable' = None,
+        params: cirq.Sweepable = None,
         permit_terminal_measurements: bool = False,
     ) -> Sequence[Sequence[float]]:
         """Calculates estimated expectation values from samples of a circuit.
@@ -360,17 +359,17 @@ class Sampler(metaclass=value.ABCMetaImplementAnyOneOf):
             )
 
         # Wrap input into a list of pauli sum
-        pauli_sums: List['cirq.PauliSum'] = (
+        pauli_sums: list[cirq.PauliSum] = (
             [ops.PauliSum.wrap(o) for o in observables]
-            if isinstance(observables, List)
+            if isinstance(observables, list)
             else [ops.PauliSum.wrap(observables)]
         )
         del observables
 
         # Flatten Pauli Sum into one big list of Pauli String
         # Keep track of which Pauli Sum each one was from.
-        flat_pstrings: List['cirq.PauliString'] = []
-        pstring_to_psum_i: Dict['cirq.PauliString', int] = {}
+        flat_pstrings: list[cirq.PauliString] = []
+        pstring_to_psum_i: dict[cirq.PauliString, int] = {}
         for psum_i, pauli_sum in enumerate(pauli_sums):
             for pstring in pauli_sum:
                 flat_pstrings.append(pstring)
@@ -378,10 +377,10 @@ class Sampler(metaclass=value.ABCMetaImplementAnyOneOf):
 
         # Flatten Circuit Sweep into one big list of Params.
         # Keep track of their indices so we can map back.
-        flat_params: List['cirq.ParamMappingType'] = [
+        flat_params: list[cirq.ParamMappingType] = [
             pr.param_dict for pr in study.to_resolvers(params)
         ]
-        circuit_param_to_sweep_i: Dict[FrozenSet[Tuple[str, Union[int, Tuple[int, int]]]], int] = {
+        circuit_param_to_sweep_i: dict[frozenset[tuple[str, int | tuple[int, int]]], int] = {
             _hashable_param(param.items()): i for i, param in enumerate(flat_params)
         }
 
@@ -399,7 +398,7 @@ class Sampler(metaclass=value.ABCMetaImplementAnyOneOf):
         # nesting structure, we place the measured values according to the back-mappings we set up
         # above. We also do the sum operation to aggregate multiple PauliString measured values
         # for a given PauliSum.
-        nested_results: List[List[float]] = [[0] * len(pauli_sums) for _ in range(len(flat_params))]
+        nested_results: list[list[float]] = [[0] * len(pauli_sums) for _ in range(len(flat_params))]
         for res in obs_meas_results:
             param_i = circuit_param_to_sweep_i[_hashable_param(res.circuit_params.items())]
             psum_i = pstring_to_psum_i[res.setting.observable]
@@ -409,8 +408,8 @@ class Sampler(metaclass=value.ABCMetaImplementAnyOneOf):
 
     @staticmethod
     def _get_measurement_shapes(
-        circuit: 'cirq.AbstractCircuit',
-    ) -> Dict[str, Tuple[int, Tuple[int, ...]]]:
+        circuit: cirq.AbstractCircuit,
+    ) -> dict[str, tuple[int, tuple[int, ...]]]:
         """Gets the shapes of measurements in the given circuit.
 
         Returns:
@@ -423,8 +422,8 @@ class Sampler(metaclass=value.ABCMetaImplementAnyOneOf):
             ValueError: if the qid_shape of different instances of the same measurement
             key disagree.
         """
-        qid_shapes: Dict[str, Tuple[int, ...]] = {}
-        num_instances: Dict[str, int] = collections.Counter()
+        qid_shapes: dict[str, tuple[int, ...]] = {}
+        num_instances: dict[str, int] = collections.Counter()
         for op in circuit.all_operations():
             key = protocols.measurement_key_name(op, default=None)
             if key is not None: