cirq-core 1.4.0.dev20240529225110__py3-none-any.whl → 1.5.0__py3-none-any.whl

cirq/value/linear_dict.py

@@ -13,8 +13,9 @@
 # limitations under the License.
 
 """Linear combination represented as mapping of things to coefficients."""
-import numbers
+
 from typing import (
+    AbstractSet,
     Any,
     Callable,
     Dict,
@@ -28,19 +29,44 @@ from typing import (
     Optional,
     overload,
     Tuple,
+    TYPE_CHECKING,
     TypeVar,
     Union,
     ValuesView,
 )
+
+import numpy as np
+import sympy
 from typing_extensions import Self
 
-Scalar = Union[complex, float, numbers.Complex]
+from cirq import protocols
+
+if TYPE_CHECKING:
+    import cirq
+
+Scalar = Union[complex, np.number]
 TVector = TypeVar('TVector')
 
 TDefault = TypeVar('TDefault')
 
 
-def _format_coefficient(format_spec: str, coefficient: Scalar) -> str:
+class _SympyPrinter(sympy.printing.str.StrPrinter):
+    def __init__(self, format_spec: str):
+        super().__init__()
+        self._format_spec = format_spec
+
+    def _print(self, expr, **kwargs):
+        if expr.is_complex:
+            coefficient = complex(expr)
+            s = _format_coefficient(self._format_spec, coefficient)
+            return s[1:-1] if s.startswith('(') else s
+        return super()._print(expr, **kwargs)
+
+
+def _format_coefficient(format_spec: str, coefficient: 'cirq.TParamValComplex') -> str:
+    if isinstance(coefficient, sympy.Basic):
+        printer = _SympyPrinter(format_spec)
+        return printer.doprint(coefficient)
     coefficient = complex(coefficient)
     real_str = f'{coefficient.real:{format_spec}}'
     imag_str = f'{coefficient.imag:{format_spec}}'
@@ -57,7 +83,7 @@ def _format_coefficient(format_spec: str, coefficient: Scalar) -> str:
     return f'({real_str}+{imag_str}j)'
 
 
-def _format_term(format_spec: str, vector: TVector, coefficient: Scalar) -> str:
+def _format_term(format_spec: str, vector: TVector, coefficient: 'cirq.TParamValComplex') -> str:
     coefficient_str = _format_coefficient(format_spec, coefficient)
     if not coefficient_str:
         return coefficient_str
@@ -67,7 +93,7 @@ def _format_term(format_spec: str, vector: TVector, coefficient: Scalar) -> str:
     return '+' + result
 
 
-def _format_terms(terms: Iterable[Tuple[TVector, Scalar]], format_spec: str):
+def _format_terms(terms: Iterable[Tuple[TVector, 'cirq.TParamValComplex']], format_spec: str):
     formatted_terms = [_format_term(format_spec, vector, coeff) for vector, coeff in terms]
     s = ''.join(formatted_terms)
     if not s:
@@ -77,7 +103,7 @@ def _format_terms(terms: Iterable[Tuple[TVector, Scalar]], format_spec: str):
     return s
 
 
-class LinearDict(Generic[TVector], MutableMapping[TVector, Scalar]):
+class LinearDict(Generic[TVector], MutableMapping[TVector, 'cirq.TParamValComplex']):
     """Represents linear combination of things.
 
     LinearDict implements the basic linear algebraic operations of vector
@@ -94,7 +120,7 @@ class LinearDict(Generic[TVector], MutableMapping[TVector, Scalar]):
 
     def __init__(
         self,
-        terms: Optional[Mapping[TVector, Scalar]] = None,
+        terms: Optional[Mapping[TVector, 'cirq.TParamValComplex']] = None,
         validator: Optional[Callable[[TVector], bool]] = None,
     ) -> None:
         """Initializes linear combination from a collection of terms.
@@ -110,13 +136,18 @@ class LinearDict(Generic[TVector], MutableMapping[TVector, Scalar]):
         """
         self._has_validator = validator is not None
         self._is_valid = validator or (lambda x: True)
-        self._terms: Dict[TVector, Scalar] = {}
+        self._terms: Dict[TVector, 'cirq.TParamValComplex'] = {}
         if terms is not None:
             self.update(terms)
 
     @classmethod
     def fromkeys(cls, vectors, coefficient=0):
-        return LinearDict(dict.fromkeys(vectors, complex(coefficient)))
+        return LinearDict(
+            dict.fromkeys(
+                vectors,
+                coefficient if isinstance(coefficient, sympy.Basic) else complex(coefficient),
+            )
+        )
 
     def _check_vector_valid(self, vector: TVector) -> None:
         if not self._is_valid(vector):
@@ -124,7 +155,11 @@ class LinearDict(Generic[TVector], MutableMapping[TVector, Scalar]):
 
     def clean(self, *, atol: float = 1e-9) -> Self:
         """Remove terms with coefficients of absolute value atol or less."""
-        negligible = [v for v, c in self._terms.items() if abs(c) <= atol]  # type: ignore[operator]
+        negligible = [
+            v
+            for v, c in self._terms.items()
+            if not isinstance(c, sympy.Basic) and abs(complex(c)) <= atol
+        ]
         for v in negligible:
             del self._terms[v]
         return self
@@ -137,40 +172,50 @@ class LinearDict(Generic[TVector], MutableMapping[TVector, Scalar]):
         snapshot = self.copy().clean(atol=0)
         return snapshot._terms.keys()
 
-    def values(self) -> ValuesView[Scalar]:
+    def values(self) -> ValuesView['cirq.TParamValComplex']:
         snapshot = self.copy().clean(atol=0)
         return snapshot._terms.values()
 
-    def items(self) -> ItemsView[TVector, Scalar]:
+    def items(self) -> ItemsView[TVector, 'cirq.TParamValComplex']:
         snapshot = self.copy().clean(atol=0)
         return snapshot._terms.items()
 
     # pylint: disable=function-redefined
     @overload
-    def update(self, other: Mapping[TVector, Scalar], **kwargs: Scalar) -> None:
+    def update(
+        self, other: Mapping[TVector, 'cirq.TParamValComplex'], **kwargs: 'cirq.TParamValComplex'
+    ) -> None:
         pass
 
     @overload
-    def update(self, other: Iterable[Tuple[TVector, Scalar]], **kwargs: Scalar) -> None:
+    def update(
+        self,
+        other: Iterable[Tuple[TVector, 'cirq.TParamValComplex']],
+        **kwargs: 'cirq.TParamValComplex',
+    ) -> None:
         pass
 
     @overload
-    def update(self, *args: Any, **kwargs: Scalar) -> None:
+    def update(self, *args: Any, **kwargs: 'cirq.TParamValComplex') -> None:
         pass
 
     def update(self, *args, **kwargs):
         terms = dict()
         terms.update(*args, **kwargs)
         for vector, coefficient in terms.items():
+            if isinstance(coefficient, sympy.Basic):
+                coefficient = sympy.simplify(coefficient)
+                if coefficient.is_complex:
+                    coefficient = complex(coefficient)
             self[vector] = coefficient
         self.clean(atol=0)
 
     @overload
-    def get(self, vector: TVector) -> Scalar:
+    def get(self, vector: TVector) -> 'cirq.TParamValComplex':
         pass
 
     @overload
-    def get(self, vector: TVector, default: TDefault) -> Union[Scalar, TDefault]:
+    def get(self, vector: TVector, default: TDefault) -> Union['cirq.TParamValComplex', TDefault]:
         pass
 
     def get(self, vector, default=0):
@@ -183,10 +228,10 @@ class LinearDict(Generic[TVector], MutableMapping[TVector, Scalar]):
     def __contains__(self, vector: Any) -> bool:
         return vector in self._terms and self._terms[vector] != 0
 
-    def __getitem__(self, vector: TVector) -> Scalar:
+    def __getitem__(self, vector: TVector) -> 'cirq.TParamValComplex':
         return self._terms.get(vector, 0)
 
-    def __setitem__(self, vector: TVector, coefficient: Scalar) -> None:
+    def __setitem__(self, vector: TVector, coefficient: 'cirq.TParamValComplex') -> None:
         self._check_vector_valid(vector)
         if coefficient != 0:
             self._terms[vector] = coefficient
@@ -234,21 +279,21 @@ class LinearDict(Generic[TVector], MutableMapping[TVector, Scalar]):
         factory = type(self)
         return factory({v: -c for v, c in self.items()})
 
-    def __imul__(self, a: Scalar) -> Self:
+    def __imul__(self, a: 'cirq.TParamValComplex') -> Self:
         for vector in self:
             self._terms[vector] *= a
         self.clean(atol=0)
         return self
 
-    def __mul__(self, a: Scalar) -> Self:
+    def __mul__(self, a: 'cirq.TParamValComplex') -> Self:
         result = self.copy()
         result *= a
-        return result
+        return result.copy()
 
-    def __rmul__(self, a: Scalar) -> Self:
+    def __rmul__(self, a: 'cirq.TParamValComplex') -> Self:
         return self.__mul__(a)
 
-    def __truediv__(self, a: Scalar) -> Self:
+    def __truediv__(self, a: 'cirq.TParamValComplex') -> Self:
         return self.__mul__(1 / a)
 
     def __bool__(self) -> bool:
@@ -318,3 +363,19 @@ class LinearDict(Generic[TVector], MutableMapping[TVector, Scalar]):
     @classmethod
     def _from_json_dict_(cls, keys, values, **kwargs):
         return cls(terms=dict(zip(keys, values)))
+
+    def _is_parameterized_(self) -> bool:
+        return any(protocols.is_parameterized(v) for v in self._terms.values())
+
+    def _parameter_names_(self) -> AbstractSet[str]:
+        return set(name for v in self._terms.values() for name in protocols.parameter_names(v))
+
+    def _resolve_parameters_(self, resolver: 'cirq.ParamResolver', recursive: bool) -> 'LinearDict':
+        result = self.copy()
+        result.update(
+            {
+                k: protocols.resolve_parameters(v, resolver, recursive)
+                for k, v in self._terms.items()
+            }
+        )
+        return result

cirq/value/linear_dict_test.py

@@ -14,6 +14,7 @@
 
 import numpy as np
 import pytest
+import sympy
 
 import cirq
 
@@ -24,6 +25,12 @@ def test_empty_init():
     assert not v
 
 
+sym = sympy.Symbol('sym')
+expr = sym * -(2 + 3j)
+symval = expr.subs({'sym': 5})
+symvalresolved = -10 - 15j
+
+
 @pytest.mark.parametrize(
     'keys, coefficient, terms_expected',
     (
@@ -31,6 +38,9 @@ def test_empty_init():
         (('X',), 2, {'X': 2}),
         (('a', 'b', 'c', 'd'), 0.5, {'a': 0.5, 'b': 0.5, 'c': 0.5, 'd': 0.5}),
         (('b', 'c', 'd', 'e'), -2j, {'b': -2j, 'c': -2j, 'd': -2j, 'e': -2j}),
+        (('b', 'c'), sym, {'b': sym, 'c': sym}),
+        (('b', 'c'), expr, {'b': expr, 'c': expr}),
+        (('b', 'c'), symval, {'b': symvalresolved, 'c': symvalresolved}),
     ),
 )
 def test_fromkeys(keys, coefficient, terms_expected):
@@ -273,6 +283,9 @@ def test_len(terms, expected_length):
         ({'X': 1}, {'Y': 2}, {'X': 1, 'Y': 2}),
         ({'X': 1}, {'X': 1}, {'X': 2}),
         ({'X': 1, 'Y': 2}, {'Y': -2}, {'X': 1}),
+        ({'X': 1}, {'X': sym}, {'X': sym + 1}),
+        ({'X': 1}, {'X': expr}, {'X': expr + 1}),
+        ({'X': 1}, {'X': symval}, {'X': symvalresolved + 1}),
     ),
 )
 def test_vector_addition(terms_1, terms_2, terms_expected):
@@ -297,6 +310,9 @@ def test_vector_addition(terms_1, terms_2, terms_expected):
         ({'X': 1}, {'X': 1}, {}),
         ({'X': 1, 'Y': 2}, {'Y': 2}, {'X': 1}),
         ({'X': 1, 'Y': 2}, {'Y': 3}, {'X': 1, 'Y': -1}),
+        ({'X': 1}, {'X': sym}, {'X': 1 - sym}),
+        ({'X': 1}, {'X': expr}, {'X': 1 - expr}),
+        ({'X': 1}, {'X': symval}, {'X': 1 - symvalresolved}),
     ),
 )
 def test_vector_subtraction(terms_1, terms_2, terms_expected):
@@ -313,7 +329,14 @@ def test_vector_subtraction(terms_1, terms_2, terms_expected):
 
 @pytest.mark.parametrize(
     'terms, terms_expected',
-    (({}, {}), ({'key': 1}, {'key': -1}), ({'1': 10, '2': -20}, {'1': -10, '2': 20})),
+    (
+        ({}, {}),
+        ({'key': 1}, {'key': -1}),
+        ({'1': 10, '2': -20}, {'1': -10, '2': 20}),
+        ({'key': sym}, {'key': -sym}),
+        ({'key': expr}, {'key': -expr}),
+        ({'key': symval}, {'key': -symvalresolved}),
+    ),
 )
 def test_vector_negation(terms, terms_expected):
     linear_dict = cirq.LinearDict(terms)
@@ -331,6 +354,12 @@ def test_vector_negation(terms, terms_expected):
         (0, {'abc': 10, 'def': 20}, {}),
         (1j, {'X': 4j}, {'X': -4}),
         (-1, {'a': 10, 'b': -20}, {'a': -10, 'b': 20}),
+        (2, {'X': sym}, {'X': 2 * sym}),
+        (2, {'X': expr}, {'X': 2 * expr}),
+        (2, {'X': symval}, {'X': 2 * symvalresolved}),
+        (sym, {'X': 2}, {'X': 2 * sym}),
+        (expr, {'X': 2}, {'X': 2 * expr}),
+        (symval, {'X': 2}, {'X': 2 * symvalresolved}),
     ),
 )
 def test_scalar_multiplication(scalar, terms, terms_expected):
@@ -350,6 +379,12 @@ def test_scalar_multiplication(scalar, terms, terms_expected):
         (2, {'X': 6, 'Y': -2}, {'X': 3, 'Y': -1}),
         (1j, {'X': 1, 'Y': 1j}, {'X': -1j, 'Y': 1}),
         (-1, {'a': 10, 'b': -20}, {'a': -10, 'b': 20}),
+        (2, {'X': sym}, {'X': 0.5 * sym}),
+        (2, {'X': expr}, {'X': 0.5 * expr}),
+        (2, {'X': symval}, {'X': 0.5 * symvalresolved}),
+        (sym, {'X': 2}, {'X': 2 / sym}),
+        (expr, {'X': 2}, {'X': 2 / expr}),
+        (symval, {'X': 2}, {'X': 2 / symvalresolved}),
     ),
 )
 def test_scalar_division(scalar, terms, terms_expected):
@@ -387,7 +422,13 @@ def test_bool(terms, bool_value):
 
 @pytest.mark.parametrize(
     'terms_1, terms_2',
-    (({}, {}), ({}, {'X': 0}), ({'X': 0.0}, {'Y': 0.0}), ({'a': 1}, {'a': 1, 'b': 0})),
+    (
+        ({}, {}),
+        ({}, {'X': 0}),
+        ({'X': 0.0}, {'Y': 0.0}),
+        ({'a': 1}, {'a': 1, 'b': 0}),
+        ({'X': sym}, {'X': sym}),
+    ),
 )
 def test_equal(terms_1, terms_2):
     linear_dict_1 = cirq.LinearDict(terms_1)
@@ -405,6 +446,7 @@ def test_equal(terms_1, terms_2):
         ({'X': 1e-12}, {'X': 0}),
         ({'X': 0.0}, {'Y': 0.1}),
         ({'X': 1}, {'X': 1, 'Z': 1e-12}),
+        ({'X': sym + 0.1}, {'X': sym}),
     ),
 )
 def test_unequal(terms_1, terms_2):
@@ -423,6 +465,7 @@ def test_unequal(terms_1, terms_2):
         ({'X': 1e-12}, {'X': 0}),
         ({'X': 5e-10}, {'Y': 2e-11}),
         ({'X': 1.000000001}, {'X': 1, 'Z': 0}),
+        ({'X': sym + 0.000000001}, {'X': sym}),
     ),
 )
 def test_approximately_equal(terms_1, terms_2):
@@ -476,7 +519,20 @@ def test_format(terms, fmt, expected_string):
     assert actual_string.replace(' ', '') == expected_string.replace(' ', '')
 
 
-@pytest.mark.parametrize('terms', (({}, {'X': 1}, {'X': 2, 'Y': 3}, {'X': 1.23456789e-12})))
+@pytest.mark.parametrize(
+    'terms',
+    (
+        (
+            {},
+            {'X': 1},
+            {'X': 2, 'Y': 3},
+            {'X': 1.23456789e-12},
+            {'X': sym},
+            ({'X': sym * 2}),
+            {'X': symval},
+        )
+    ),
+)
 def test_repr(terms):
     original = cirq.LinearDict(terms)
     recovered = eval(repr(original))
@@ -502,6 +558,10 @@ def test_repr(terms):
         ({'X': -2, 'Y': -3}, '-2.000*X-3.000*Y'),
         ({'X': -2j, 'Y': -3}, '-2.000j*X-3.000*Y'),
         ({'X': -2j, 'Y': -3j}, '-2.000j*X-3.000j*Y'),
+        ({'X': sym}, 'sym*X'),
+        ({'X': sym * 2}, '2.000*sym*X'),
+        ({'X': expr}, '-sym*(2.000+3.000j)*X'),
+        ({'X': symval}, '-(10.000+15.000j)*X'),
     ),
 )
 def test_str(terms, string):
@@ -547,3 +607,34 @@ def test_repr_pretty(terms):
 def test_json_fails_with_validator():
     with pytest.raises(ValueError, match='not json serializable'):
         _ = cirq.to_json(cirq.LinearDict({}, validator=lambda: True))
+
+
+@pytest.mark.parametrize(
+    'terms, names',
+    (
+        ({'X': sym}, {'sym'}),
+        ({'X': sym * sympy.Symbol('a')}, {'sym', 'a'}),
+        ({'X': expr}, {'sym'}),
+        ({'X': sym, 'Y': sympy.Symbol('a')}, {'sym', 'a'}),
+        ({'X': symval}, set()),
+    ),
+)
+def test_parameter_names(terms, names):
+    linear_dict = cirq.LinearDict(terms)
+    assert cirq.parameter_names(linear_dict) == names
+
+
+@pytest.mark.parametrize(
+    'terms, expected',
+    (
+        ({'X': sym}, {'X': 2}),
+        ({'X': sym * sympy.Symbol('a')}, {'X': 6}),
+        ({'X': expr}, {'X': -4 - 6j}),
+        ({'X': sym, 'Y': sympy.Symbol('a')}, {'X': 2, 'Y': 3}),
+        ({'X': symval}, {'X': symvalresolved}),
+    ),
+)
+def test_resolve_parameters(terms, expected):
+    linear_dict = cirq.LinearDict(terms)
+    expected_dict = cirq.LinearDict(expected)
+    assert cirq.resolve_parameters(linear_dict, {'sym': 2, 'a': 3}) == expected_dict

cirq/value/measurement_key.py

@@ -12,9 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import FrozenSet, Mapping, Optional, Tuple
-
 import dataclasses
+from typing import Any, Dict, FrozenSet, Mapping, Optional, Tuple
 
 MEASUREMENT_KEY_SEPARATOR = ':'
 
@@ -77,6 +76,14 @@ class MeasurementKey:
             object.__setattr__(self, '_hash', hash(str(self)))
         return self._hash
 
+    def __getstate__(self) -> Dict[str, Any]:
+        # clear cached hash value when pickling, see #6674
+        state = self.__dict__
+        if "_hash" in state:
+            state = state.copy()
+            del state["_hash"]
+        return state
+
     def __lt__(self, other):
         if isinstance(other, MeasurementKey):
             if self.path != other.path:

cirq/value/periodic_value.py

@@ -12,13 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import AbstractSet, Any, TYPE_CHECKING, Union
+from typing import AbstractSet, Any, TYPE_CHECKING
 
 import sympy
 
 from cirq._compat import proper_repr
 
-
 if TYPE_CHECKING:
     import cirq
 
@@ -36,7 +35,7 @@ class PeriodicValue:
     interval.
     """
 
-    def __init__(self, value: Union[int, float, sympy.Expr], period: Union[int, float, sympy.Expr]):
+    def __init__(self, value: 'cirq.TParamVal', period: 'cirq.TParamVal'):
         """Initializes the equivalence class.
 
         Args:

cirq/value/periodic_value_test.py

@@ -42,6 +42,11 @@ def test_periodic_value_approx_eq_basic():
     assert not cirq.approx_eq(cirq.PeriodicValue(1.0, 2.0), cirq.PeriodicValue(1.0, 2.2), atol=0.1)
     assert not cirq.approx_eq(cirq.PeriodicValue(1.0, 2.0), cirq.PeriodicValue(1.2, 2.2), atol=0.3)
     assert not cirq.approx_eq(cirq.PeriodicValue(1.0, 2.0), cirq.PeriodicValue(1.2, 2.2), atol=0.1)
+    assert cirq.approx_eq(
+        cirq.PeriodicValue(sympy.Symbol('t'), 2.0),
+        cirq.PeriodicValue(sympy.Symbol('t'), 2.0),
+        atol=0.1,
+    )
 
 
 def test_periodic_value_approx_eq_normalized():

cirq/value/probability.py

@@ -17,6 +17,7 @@
 from typing import TYPE_CHECKING
 
 import numpy as np
+
 from cirq.qis import to_valid_state_vector
 
 if TYPE_CHECKING:

cirq/value/random_state.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import cast, Any
+from typing import Any, cast
 
 import numpy as np
 

cirq/value/timestamp.py

@@ -14,7 +14,7 @@
 """A typed location in time that supports picosecond accuracy."""
 
 from datetime import timedelta
-from typing import Union, overload
+from typing import overload
 
 from cirq.value.duration import Duration
 
@@ -24,9 +24,7 @@ class Timestamp:
 
     Supports affine operations against Duration."""
 
-    def __init__(
-        self, *, picos: Union[int, float] = 0, nanos: Union[int, float] = 0  # Forces keyword args.
-    ) -> None:
+    def __init__(self, *, picos: float = 0, nanos: float = 0) -> None:  # Forces keyword args.
         """Initializes a Timestamp with a time specified in ns and/or ps.
 
         The time is relative to some unspecified "time zero". If both picos and

cirq/value/timestamp_test.py

@@ -13,10 +13,11 @@
 # limitations under the License.
 
 from datetime import timedelta
+
 import pytest
 
 import cirq
-from cirq import Timestamp, Duration
+from cirq import Duration, Timestamp
 
 
 def test_init():

cirq/value/type_alias.py

@@ -14,10 +14,10 @@
 
 from typing import Union
 
+import numpy as np
 import sympy
 
 from cirq._doc import document
-from cirq.value import linear_dict
 
 """Supply aliases for commonly used types.
 """
@@ -28,5 +28,5 @@ document(TParamKey, """A parameter that a parameter resolver may map to a value.
 TParamVal = Union[float, sympy.Expr]
 document(TParamVal, """A value that a parameter resolver may return for a parameter.""")
 
-TParamValComplex = Union[linear_dict.Scalar, sympy.Expr]
+TParamValComplex = Union[complex, np.number, sympy.Expr]
 document(TParamValComplex, """A complex value that parameter resolvers may use for parameters.""")

cirq/value/value_equality_attr.py

@@ -13,11 +13,11 @@
 # limitations under the License.
 """Defines `@cirq.value_equality`, for easy __eq__/__hash__ methods."""
 
-from typing import Any, Callable, Optional, overload, Union
+from typing import Any, Callable, Dict, Optional, overload, Union
 
 from typing_extensions import Protocol
 
-from cirq import protocols, _compat
+from cirq import _compat, protocols
 
 
 class _SupportsValueEquality(Protocol):
@@ -110,6 +110,16 @@ def _value_equality_approx_eq(
     )
 
 
+def _value_equality_getstate(self: _SupportsValueEquality) -> Dict[str, Any]:
+    # clear cached hash value when pickling, see #6674
+    state = self.__dict__
+    hash_attr = _compat._method_cache_name(self.__hash__)
+    if hash_attr in state:
+        state = state.copy()
+        del state[hash_attr]
+    return state
+
+
 # pylint: disable=function-redefined
 @overload
 def value_equality(
@@ -228,6 +238,8 @@ def value_equality(
     cached_values_getter = values_getter if unhashable else _compat.cached_method(values_getter)
     setattr(cls, '_value_equality_values_', cached_values_getter)
     setattr(cls, '__hash__', None if unhashable else _compat.cached_method(_value_equality_hash))
+    if not unhashable:
+        setattr(cls, '__getstate__', _value_equality_getstate)
     setattr(cls, '__eq__', _value_equality_eq)
     setattr(cls, '__ne__', _value_equality_ne)
 

cirq/value/value_equality_attr_test.py

@@ -94,7 +94,7 @@ class UnhashableD:
         self.x = x
 
     def _value_equality_values_(self):
-        return self.x
+        return self.x  # pragma: no cover
 
 
 class UnhashableCa(UnhashableC):

cirq/vis/__init__.py

@@ -14,14 +14,17 @@
 
 """Classes and methods for visualizing results, quantum states, and devices."""
 
-from cirq.vis.heatmap import Heatmap
+from cirq.vis.heatmap import Heatmap as Heatmap
 
-from cirq.vis.heatmap import TwoQubitInteractionHeatmap
+from cirq.vis.heatmap import TwoQubitInteractionHeatmap as TwoQubitInteractionHeatmap
 
-from cirq.vis.histogram import integrated_histogram
+from cirq.vis.histogram import integrated_histogram as integrated_histogram
 
-from cirq.vis.state_histogram import get_state_histogram, plot_state_histogram
+from cirq.vis.state_histogram import (
+    get_state_histogram as get_state_histogram,
+    plot_state_histogram as plot_state_histogram,
+)
 
-from cirq.vis.density_matrix import plot_density_matrix
+from cirq.vis.density_matrix import plot_density_matrix as plot_density_matrix
 
-from cirq.vis.vis_utils import relative_luminance
+from cirq.vis.vis_utils import relative_luminance as relative_luminance

cirq/vis/density_matrix.py

@@ -16,8 +16,8 @@
 
 from typing import Optional
 
-import numpy as np
 import matplotlib.pyplot as plt
+import numpy as np
 from matplotlib import lines, patches
 
 from cirq.qis.states import validate_density_matrix

cirq/vis/density_matrix_test.py

@@ -15,13 +15,10 @@
 
 import numpy as np
 import pytest
-
-from matplotlib import lines, patches, text, spines, axis
-from matplotlib import pyplot as plt
+from matplotlib import axis, lines, patches, pyplot as plt, spines, text
 
 import cirq.testing
-from cirq.vis.density_matrix import plot_density_matrix
-from cirq.vis.density_matrix import _plot_element_of_density_matrix
+from cirq.vis.density_matrix import _plot_element_of_density_matrix, plot_density_matrix
 
 
 @pytest.mark.usefixtures('closefigures')