From e18824736d40a62376aad220b502419a68ffd572 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Wed, 17 Jul 2024 17:00:09 -0500 Subject: [PATCH 1/9] change attrs to dataclasses --- doc/conf.py | 1 - doc/design.rst | 17 ++-- pytato/array.py | 97 ++++++++++----------- pytato/distributed/nodes.py | 14 +-- pytato/distributed/partition.py | 10 +-- pytato/distributed/tags.py | 2 +- pytato/distributed/verify.py | 10 +-- pytato/function.py | 26 +++--- pytato/loopy.py | 13 ++- pytato/scalar_expr.py | 8 +- pytato/stringifier.py | 8 +- pytato/target/loopy/codegen.py | 21 +++-- pytato/transform/__init__.py | 8 +- pytato/transform/einsum_distributive_law.py | 14 +-- pytato/visualization/dot.py | 8 +- setup.py | 5 +- test/test_pytato.py | 6 +- 17 files changed, 123 insertions(+), 145 deletions(-) diff --git a/doc/conf.py b/doc/conf.py index 75f5beab7..091e30e2f 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -27,7 +27,6 @@ "sumpy": ("https://documen.tician.de/sumpy/", None), "islpy": ("https://documen.tician.de/islpy/", None), "jax": ("https://jax.readthedocs.io/en/latest/", None), - "attrs": ("https://www.attrs.org/en/stable/", None), "mpi4py": ("https://mpi4py.readthedocs.io/en/latest", None), "immutabledict": ("https://immutabledict.corenting.fr/", None), } diff --git a/doc/design.rst b/doc/design.rst index 93a4ebd23..d7791cab7 100644 --- a/doc/design.rst +++ b/doc/design.rst @@ -226,20 +226,15 @@ that relies on memory layout information to do its job is undefined in :mod:`pyt At the most basic level, the attribute :attr:`numpy.ndarray.strides` is not available on subclasses of :class:`pytato.Array`. -Dataclasses / :mod:`attrs` --------------------------- +Dataclasses +----------- :mod:`dataclasses` helps us reduce most of the boilerplate involved in -instantiating a new type. However, :mod:`dataclasses` does not support -keyword-only argument until Python-3.10. To overcome this, we prefer -:mod:`attrs` which gives us all the required functionality of -:mod:`dataclasses` and works with Python-3.8. - - +instantiating a new type. We have checks in place to avoid developer errors that could happen by using -the defaults of these libraries. For eg. both :mod:`dataclasses` and -:mod:`attrs` override the implementation of ``__eq__`` for the class being -implemented, which could potentially lead lead to an `exponential complex +the defaults of this library. For example, :mod:`dataclasses` overrides the +implementation of ``__eq__`` for the class being implemented, which could +potentially lead to an `exponentially complex operation `_. Lessons learned diff --git a/pytato/array.py b/pytato/array.py index 991eb11f5..bfc47d0af 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -172,7 +172,7 @@ from enum import IntEnum from functools import partialmethod, cached_property import operator -import attrs +import dataclasses from typing import ( Optional, Callable, ClassVar, Dict, Any, Mapping, Tuple, Union, Protocol, Sequence, cast, TYPE_CHECKING, List, Iterator, TypeVar, @@ -379,7 +379,7 @@ def _truediv_result_type(*dtypes: DtypeOrPyScalarType) -> np.dtype[Any]: return dtype -@attrs.frozen +@dataclasses.dataclass(frozen=True) class NormalizedSlice: """ A normalized version of :class:`slice`. "Normalized" is explained in @@ -404,7 +404,7 @@ class NormalizedSlice: step: IntegralT -@attrs.frozen +@dataclasses.dataclass(frozen=True) class Axis(Taggable): """ A type for recording the information about an :class:`~pytato.Array`'s @@ -413,11 +413,11 @@ class Axis(Taggable): tags: FrozenSet[Tag] def _with_new_tags(self, tags: FrozenSet[Tag]) -> Axis: - from attrs import evolve as replace + from dataclasses import replace return replace(self, tags=tags) -@attrs.frozen +@dataclasses.dataclass(frozen=True) class ReductionDescriptor(Taggable): """ Records information about a reduction dimension in an @@ -426,11 +426,11 @@ class ReductionDescriptor(Taggable): tags: FrozenSet[Tag] def _with_new_tags(self, tags: FrozenSet[Tag]) -> ReductionDescriptor: - from attrs import evolve as replace + from dataclasses import replace return replace(self, tags=tags) -@attrs.frozen(eq=False, repr=False, hash=True, cache_hash=True) +@dataclasses.dataclass(frozen=True, eq=False, repr=False, unsafe_hash=True) class Array(Taggable): r""" A base class (abstract interface + supplemental functionality) for lazily @@ -520,11 +520,11 @@ class Array(Taggable): .. attribute:: ndim """ - axes: AxesT = attrs.field(kw_only=True) - tags: FrozenSet[Tag] = attrs.field(kw_only=True) + axes: AxesT = dataclasses.field(kw_only=True) + tags: FrozenSet[Tag] = dataclasses.field(kw_only=True) # These are automatically excluded from equality in EqualityComparer - non_equality_tags: FrozenSet[Tag] = attrs.field(kw_only=True, + non_equality_tags: FrozenSet[Tag] = dataclasses.field(kw_only=True, hash=False, default=frozenset()) @@ -537,7 +537,7 @@ def _is_eq_valid(self) -> bool: return self.__class__.__eq__ is Array.__eq__ if __debug__: - def __attrs_post_init__(self) -> None: + def __post_init__(self) -> None: if _ENABLE_TRACEBACK_TAG: from pytato.tags import CreatedAt ntags = sum( @@ -553,10 +553,10 @@ def __attrs_post_init__(self) -> None: assert self._is_eq_valid() def copy(self: ArrayT, **kwargs: Any) -> ArrayT: - return attrs.evolve(self, **kwargs) + return dataclasses.replace(self, **kwargs) def _with_new_tags(self: ArrayT, tags: FrozenSet[Tag]) -> ArrayT: - return attrs.evolve(self, tags=tags) + return dataclasses.replace(self, tags=tags) if TYPE_CHECKING: @property @@ -786,7 +786,7 @@ def __repr__(self) -> str: # {{{ mixins -@attrs.frozen(eq=False, slots=False, repr=False) +@dataclasses.dataclass(frozen=True, eq=False, slots=False, repr=False) class _SuppliedShapeAndDtypeMixin: """A mixin class for when an array must store its own *shape* and *dtype*, rather than when it can derive them easily from inputs. @@ -799,7 +799,7 @@ class _SuppliedShapeAndDtypeMixin: # {{{ dict of named arrays -@attrs.frozen(eq=False, repr=False, hash=True, cache_hash=True) +@dataclasses.dataclass(frozen=True, eq=False, repr=False, unsafe_hash=True) class NamedArray(Array): """An entry in a :class:`AbstractResultWithNamedArrays`. Holds a reference back to the containing instance as well as the name by which *self* is @@ -850,7 +850,7 @@ def dtype(self) -> np.dtype[Any]: return self.expr.dtype -@attrs.frozen(eq=False, hash=True, cache_hash=True) +@dataclasses.dataclass(frozen=True, eq=False, unsafe_hash=True) class AbstractResultWithNamedArrays(Mapping[str, NamedArray], Taggable, ABC): r"""An abstract array computation that results in multiple :class:`Array`\ s, each named. The way in which the values of these arrays are computed @@ -866,14 +866,14 @@ class AbstractResultWithNamedArrays(Mapping[str, NamedArray], Taggable, ABC): This container deliberately does not implement arithmetic. """ - tags: FrozenSet[Tag] = attrs.field(kw_only=True) + tags: FrozenSet[Tag] = dataclasses.field(kw_only=True) _mapper_method: ClassVar[str] def _is_eq_valid(self) -> bool: return self.__class__.__eq__ is AbstractResultWithNamedArrays.__eq__ if __debug__: - def __attrs_post_init__(self) -> None: + def __post_init__(self) -> None: # ensure that a developer does not uses dataclass' "__eq__" # or "__hash__" implementation as they have exponential complexity. assert self._is_eq_valid() @@ -898,7 +898,7 @@ def __eq__(self, other: Any) -> bool: return EqualityComparer()(self, other) -@attrs.frozen(eq=False, init=False) +@dataclasses.dataclass(frozen=True, eq=False, init=False) class DictOfNamedArrays(AbstractResultWithNamedArrays): """A container of named results, each of which can be computed as an array expression provided to the constructor. @@ -907,8 +907,7 @@ class DictOfNamedArrays(AbstractResultWithNamedArrays): .. automethod:: __init__ """ - _data: Mapping[str, Array] = attrs.field( - validator=attrs.validators.instance_of(immutabledict)) + _data: Mapping[str, Array] _mapper_method: ClassVar[str] = "map_dict_of_named_arrays" @@ -954,7 +953,7 @@ def __repr__(self) -> str: # {{{ index lambda -@attrs.frozen(eq=False, repr=False, hash=True, cache_hash=True) +@dataclasses.dataclass(frozen=True, eq=False, repr=False, unsafe_hash=True) class IndexLambda(_SuppliedShapeAndDtypeMixin, Array): r"""Represents an array that can be computed by evaluating :attr:`expr` for every value of the input indices. The @@ -990,10 +989,8 @@ class IndexLambda(_SuppliedShapeAndDtypeMixin, Array): .. automethod:: with_tagged_reduction """ expr: prim.Expression - bindings: Mapping[str, Array] = attrs.field( - validator=attrs.validators.instance_of(immutabledict)) - var_to_reduction_descr: Mapping[str, ReductionDescriptor] = \ - attrs.field(validator=attrs.validators.instance_of(immutabledict)) + bindings: Mapping[str, Array] + var_to_reduction_descr: Mapping[str, ReductionDescriptor] _mapper_method: ClassVar[str] = "map_index_lambda" @@ -1049,7 +1046,7 @@ class EinsumAxisDescriptor: pass -@attrs.frozen(order=True) +@dataclasses.dataclass(frozen=True, order=True) class EinsumElementwiseAxis(EinsumAxisDescriptor): """ Describes an elementwise access pattern of an array's axis. In terms of the @@ -1059,7 +1056,7 @@ class EinsumElementwiseAxis(EinsumAxisDescriptor): dim: int -@attrs.frozen(order=True) +@dataclasses.dataclass(frozen=True, order=True) class EinsumReductionAxis(EinsumAxisDescriptor): """ Describes a reduction access pattern of an array's axis. In terms of the @@ -1069,7 +1066,7 @@ class EinsumReductionAxis(EinsumAxisDescriptor): dim: int -@attrs.frozen(frozen=True, eq=False, repr=False, hash=True, cache_hash=True) +@dataclasses.dataclass(frozen=True, eq=False, repr=False, unsafe_hash=True) class Einsum(Array): """ An array expression using the `Einstein summation convention @@ -1106,10 +1103,8 @@ class Einsum(Array): access_descriptors: Tuple[Tuple[EinsumAxisDescriptor, ...], ...] args: Tuple[Array, ...] redn_axis_to_redn_descr: Mapping[EinsumReductionAxis, - ReductionDescriptor] = \ - attrs.field(validator=attrs.validators.instance_of(immutabledict)) - index_to_access_descr: Mapping[str, EinsumAxisDescriptor] = \ - attrs.field(validator=attrs.validators.instance_of(immutabledict)) + ReductionDescriptor] + index_to_access_descr: Mapping[str, EinsumAxisDescriptor] _mapper_method: ClassVar[str] = "map_einsum" @memoize_method @@ -1421,7 +1416,7 @@ def einsum(subscripts: str, *operands: Array, # {{{ stack -@attrs.frozen(eq=False, repr=False, hash=True, cache_hash=True) +@dataclasses.dataclass(frozen=True, eq=False, repr=False, unsafe_hash=True) class Stack(Array): """Join a sequence of arrays along a new axis. @@ -1454,7 +1449,7 @@ def shape(self) -> ShapeType: # {{{ concatenate -@attrs.frozen(eq=False, repr=False, hash=True, cache_hash=True) +@dataclasses.dataclass(frozen=True, eq=False, repr=False, unsafe_hash=True) class Concatenate(Array): """Join a sequence of arrays along an existing axis. @@ -1491,7 +1486,7 @@ def shape(self) -> ShapeType: # {{{ index remapping -@attrs.frozen(eq=False, repr=False, hash=True, cache_hash=True) +@dataclasses.dataclass(frozen=True, eq=False, repr=False, unsafe_hash=True) class IndexRemappingBase(Array): """Base class for operations that remap the indices of an array. @@ -1514,7 +1509,7 @@ def dtype(self) -> np.dtype[Any]: # {{{ roll -@attrs.frozen(eq=False, repr=False, hash=True, cache_hash=True) +@dataclasses.dataclass(frozen=True, eq=False, repr=False, unsafe_hash=True) class Roll(IndexRemappingBase): """Roll an array along an axis. @@ -1540,7 +1535,7 @@ def shape(self) -> ShapeType: # {{{ axis permutation -@attrs.frozen(eq=False, repr=False, hash=True, cache_hash=True) +@dataclasses.dataclass(frozen=True, eq=False, repr=False, unsafe_hash=True) class AxisPermutation(IndexRemappingBase): r"""Permute the axes of an array. @@ -1567,7 +1562,7 @@ def shape(self) -> ShapeType: # {{{ reshape -@attrs.frozen(eq=False, repr=False, hash=True, cache_hash=True) +@dataclasses.dataclass(frozen=True, eq=False, repr=False, unsafe_hash=True) class Reshape(IndexRemappingBase): """ Reshape an array. @@ -1590,10 +1585,10 @@ class Reshape(IndexRemappingBase): _mapper_method: ClassVar[str] = "map_reshape" if __debug__: - def __attrs_post_init__(self) -> None: + def __post_init__(self) -> None: # FIXME: Get rid of this restriction assert self.order == "C" - super().__attrs_post_init__() + super().__post_init__() @property def shape(self) -> ShapeType: @@ -1604,7 +1599,7 @@ def shape(self) -> ShapeType: # {{{ indexing -@attrs.frozen(eq=False, repr=False, hash=True, cache_hash=True) +@dataclasses.dataclass(frozen=True, eq=False, repr=False, unsafe_hash=True) class IndexBase(IndexRemappingBase): """ Abstract class for all index expressions on an array. @@ -1718,7 +1713,7 @@ def shape(self) -> ShapeType: # {{{ base class for arguments -@attrs.frozen(eq=False, repr=False, hash=True, cache_hash=True) +@dataclasses.dataclass(frozen=True, eq=False, repr=False, unsafe_hash=True) class InputArgumentBase(Array): r"""Base class for input arguments. @@ -1758,7 +1753,7 @@ def dtype(self) -> np.dtype[Any]: pass -@attrs.frozen(eq=False, repr=False, hash=False) +@dataclasses.dataclass(frozen=True, eq=False, repr=False) class DataWrapper(InputArgumentBase): """Takes concrete array data and packages it to be compatible with the :class:`Array` interface. @@ -1799,7 +1794,7 @@ class DataWrapper(InputArgumentBase): (i.e. the very same instance). """ data: DataInterface - _shape: ShapeType + shape: ShapeType _mapper_method: ClassVar[str] = "map_data_wrapper" @@ -1818,10 +1813,6 @@ def __hash__(self) -> int: # it. return id(self) - @property - def shape(self) -> ShapeType: - return self._shape - @property def dtype(self) -> np.dtype[Any]: return self.data.dtype @@ -1831,7 +1822,7 @@ def dtype(self) -> np.dtype[Any]: # {{{ placeholder -@attrs.frozen(eq=False, repr=False, hash=True, cache_hash=True) +@dataclasses.dataclass(frozen=True, eq=False, repr=False, unsafe_hash=True) class Placeholder(_SuppliedShapeAndDtypeMixin, InputArgumentBase): r"""A named placeholder for an array whose concrete value is supplied by the user during evaluation. @@ -1852,7 +1843,7 @@ class Placeholder(_SuppliedShapeAndDtypeMixin, InputArgumentBase): # {{{ size parameter -@attrs.frozen(eq=False, repr=False, hash=True, cache_hash=True) +@dataclasses.dataclass(frozen=True, eq=False, repr=False, unsafe_hash=True) class SizeParam(InputArgumentBase): r"""A named placeholder for a scalar that may be used as a variable in symbolic expressions for array sizes. @@ -1863,7 +1854,7 @@ class SizeParam(InputArgumentBase): begins. """ name: str - axes: AxesT = attrs.field(kw_only=True, default=()) + axes: AxesT = dataclasses.field(kw_only=True, default=()) _mapper_method: ClassVar[str] = "map_size_param" @@ -2327,7 +2318,7 @@ def eye(N: int, M: Optional[int] = None, k: int = 0, # noqa: N803 # {{{ arange -@attrs.define +@dataclasses.dataclass class _ArangeInfo: start: Optional[int] stop: Optional[int] diff --git a/pytato/distributed/nodes.py b/pytato/distributed/nodes.py index ec2d17417..450843ed3 100644 --- a/pytato/distributed/nodes.py +++ b/pytato/distributed/nodes.py @@ -54,7 +54,7 @@ from typing import Hashable, FrozenSet, Optional, Any, ClassVar -import attrs +import dataclasses import numpy as np from pytools.tag import Taggable, Tag @@ -69,7 +69,7 @@ # {{{ send -@attrs.frozen(init=True, eq=True, hash=True, cache_hash=True) +@dataclasses.dataclass(init=True, frozen=True) class DistributedSend(Taggable): """Class representing a distributed send operation. See :class:`DistributedSendRefHolder` for a way to ensure that nodes @@ -93,20 +93,20 @@ class DistributedSend(Taggable): data: Array dest_rank: int comm_tag: CommTagType - tags: FrozenSet[Tag] = attrs.field(kw_only=True, default=frozenset()) + tags: FrozenSet[Tag] = dataclasses.field(kw_only=True, default=frozenset()) def _with_new_tags(self, tags: FrozenSet[Tag]) -> DistributedSend: - return attrs.evolve(self, tags=tags) + return dataclasses.replace(self, tags=tags) def copy(self, **kwargs: Any) -> DistributedSend: - return attrs.evolve(self, **kwargs) + return dataclasses.replace(self, **kwargs) # }}} # {{{ send ref holder -@attrs.frozen(eq=False, repr=False, init=False, hash=True) +@dataclasses.dataclass(eq=False, frozen=True, repr=False, init=False, unsafe_hash=True) class DistributedSendRefHolder(Array): """A node acting as an identity on :attr:`passthrough_data` while also holding a reference to a :class:`DistributedSend` in :attr:`send`. Since @@ -183,7 +183,7 @@ def copy(self, **kwargs: Any) -> DistributedSendRefHolder: # {{{ receive -@attrs.frozen(eq=False, hash=True, cache_hash=True) +@dataclasses.dataclass(frozen=True, eq=False, unsafe_hash=True) class DistributedRecv(_SuppliedShapeAndDtypeMixin, Array): """Class representing a distributed receive operation. diff --git a/pytato/distributed/partition.py b/pytato/distributed/partition.py index c0ba404ff..b55592ec3 100644 --- a/pytato/distributed/partition.py +++ b/pytato/distributed/partition.py @@ -67,7 +67,7 @@ Iterator, Iterable, Sequence, Any, Mapping, FrozenSet, Set, Dict, cast, List, AbstractSet, TypeVar, TYPE_CHECKING, Hashable, Optional, Tuple) -import attrs +import dataclasses from immutabledict import immutabledict from pytools.graph import CycleError @@ -92,7 +92,7 @@ import mpi4py.MPI -@attrs.define(frozen=True) +@dataclasses.dataclass(frozen=True) class CommunicationOpIdentifier: """Identifies a communication operation (consisting of a pair of a send and a receive). @@ -181,7 +181,7 @@ def __sub__(self, other: AbstractSet[_ValueT]) -> _OrderedSet[_ValueT]: PartId = Hashable -@attrs.define(frozen=True, slots=False) +@dataclasses.dataclass(frozen=True, slots=False) class DistributedGraphPart: """For one graph part, record send/receive information for input/ output names. @@ -241,7 +241,7 @@ def all_input_names(self) -> FrozenSet[str]: # {{{ distributed graph partition -@attrs.define(frozen=True, slots=False) +@dataclasses.dataclass(frozen=True, slots=False) class DistributedGraphPartition: """ .. attribute:: parts @@ -360,7 +360,7 @@ def _get_placeholder_for(self, name: str, expr: Array) -> Placeholder: # }}} -@attrs.define(frozen=True) +@dataclasses.dataclass(frozen=True) class _PartCommIDs: """A *part*, unlike a *batch*, begins with receives and ends with sends. """ diff --git a/pytato/distributed/tags.py b/pytato/distributed/tags.py index 4e33055a9..dbda66569 100644 --- a/pytato/distributed/tags.py +++ b/pytato/distributed/tags.py @@ -101,7 +101,7 @@ def number_distributed_tags( else: sym_tag_to_int_tag, next_tag = mpi_communicator.bcast(None, root=root_rank) - from attrs import evolve as replace + from dataclasses import replace return DistributedGraphPartition( parts={ pid: replace(part, diff --git a/pytato/distributed/verify.py b/pytato/distributed/verify.py index ea905357f..32d4b6c4d 100644 --- a/pytato/distributed/verify.py +++ b/pytato/distributed/verify.py @@ -37,7 +37,7 @@ from typing import Any, List, FrozenSet, Dict, Set, Optional, Sequence, TYPE_CHECKING -import attrs +import dataclasses import numpy as np from pymbolic.mapper.optimize import optimize_mapper @@ -60,7 +60,7 @@ # {{{ data structures -@attrs.define(frozen=True) +@dataclasses.dataclass(frozen=True) class _SummarizedDistributedSend: src_rank: int dest_rank: int @@ -70,19 +70,19 @@ class _SummarizedDistributedSend: dtype: np.dtype[Any] -@attrs.define(frozen=True) +@dataclasses.dataclass(frozen=True) class _DistributedPartId: rank: int part_id: PartId -@attrs.define(frozen=True) +@dataclasses.dataclass(frozen=True) class _DistributedName: rank: int name: str -@attrs.define(frozen=True) +@dataclasses.dataclass(frozen=True) class _SummarizedDistributedGraphPart: pid: _DistributedPartId needed_pids: FrozenSet[_DistributedPartId] diff --git a/pytato/function.py b/pytato/function.py index 2b492583e..0b8630578 100644 --- a/pytato/function.py +++ b/pytato/function.py @@ -43,7 +43,7 @@ THE SOFTWARE. """ -import attrs +import dataclasses import re import enum @@ -71,8 +71,8 @@ class ReturnType(enum.Enum): TUPLE_OF_ARRAYS = 2 -# eq=False to avoid equality comparison without EqualityMaper -@attrs.define(frozen=True, eq=False, hash=True) +# eq=False to avoid equality comparison without EqualityMapper +@dataclasses.dataclass(frozen=True, eq=False, unsafe_hash=True) class FunctionDefinition(Taggable): r""" A function definition that represents its outputs as instances of @@ -125,9 +125,8 @@ class FunctionDefinition(Taggable): """ parameters: FrozenSet[str] return_type: ReturnType - returns: Mapping[str, Array] = attrs.field( - validator=attrs.validators.instance_of(immutabledict)) - tags: FrozenSet[Tag] = attrs.field(kw_only=True) + returns: Mapping[str, Array] + tags: FrozenSet[Tag] = dataclasses.field(kw_only=True) @cached_property def _placeholders(self) -> Mapping[str, Placeholder]: @@ -161,7 +160,7 @@ def get_placeholder(self, name: str) -> Placeholder: def _with_new_tags( self: FunctionDefinition, tags: FrozenSet[Tag]) -> FunctionDefinition: - return attrs.evolve(self, tags=tags) + return dataclasses.replace(self, tags=tags) def __call__(self, **kwargs: Array ) -> Union[Array, @@ -263,7 +262,7 @@ def dtype(self) -> _dtype_any: # eq=False to avoid equality comparison without EqualityMapper -@attrs.define(frozen=True, eq=False, hash=True, cache_hash=True, repr=False) +@dataclasses.dataclass(frozen=True, eq=False, unsafe_hash=True, repr=False) class Call(AbstractResultWithNamedArrays): """ Records an invocation to a :class:`FunctionDefinition`. @@ -279,19 +278,18 @@ class Call(AbstractResultWithNamedArrays): """ function: FunctionDefinition - bindings: Mapping[str, Array] = attrs.field( - validator=attrs.validators.instance_of(immutabledict)) + bindings: Mapping[str, Array] _mapper_method: ClassVar[str] = "map_call" - copy = attrs.evolve + copy = dataclasses.replace if __debug__: - def __attrs_post_init__(self) -> None: + def __post_init__(self) -> None: # check that the invocation parameters and the function definition # parameters agree with each other. assert frozenset(self.bindings) == self.function.parameters - super().__attrs_post_init__() + super().__post_init__() def __contains__(self, name: object) -> bool: return name in self.function.returns @@ -306,7 +304,7 @@ def __len__(self) -> int: return len(self.function.returns) def _with_new_tags(self: Call, tags: FrozenSet[Tag]) -> Call: - return attrs.evolve(self, tags=tags) + return dataclasses.replace(self, tags=tags) # }}} diff --git a/pytato/loopy.py b/pytato/loopy.py index 0c6266cd5..e84199354 100644 --- a/pytato/loopy.py +++ b/pytato/loopy.py @@ -26,7 +26,7 @@ import numpy as np -import attrs +import dataclasses import loopy as lp import pymbolic.primitives as prim from typing import (Dict, Optional, Any, Iterator, FrozenSet, Union, Sequence, @@ -71,20 +71,19 @@ """ -@attrs.frozen(eq=False) +@dataclasses.dataclass(frozen=True, eq=False) class LoopyCall(AbstractResultWithNamedArrays): """ An array expression node representing a call to an entrypoint in a :mod:`loopy` translation unit. """ translation_unit: "lp.TranslationUnit" - bindings: Mapping[str, ArrayOrScalar] = \ - attrs.field(validator=attrs.validators.instance_of(immutabledict)) + bindings: Mapping[str, ArrayOrScalar] entrypoint: str _mapper_method: ClassVar[str] = "map_loopy_call" - copy = attrs.evolve + copy = dataclasses.replace @property def _result_names(self) -> FrozenSet[str]: @@ -116,7 +115,7 @@ def __getitem__(self, name: str) -> LoopyCallResult: raise KeyError(name) # TODO: Attach a filtered set of tags from loopy's arg. - return LoopyCallResult(container=self, + return LoopyCallResult(_container=self, name=name, axes=_get_default_axes(len(self ._entry_kernel @@ -131,7 +130,7 @@ def __iter__(self) -> Iterator[str]: return iter(self._result_names) -@attrs.frozen(eq=False, hash=True, cache_hash=True) +@dataclasses.dataclass(frozen=True, eq=False) class LoopyCallResult(NamedArray): """ Named array for :class:`LoopyCall`'s result. diff --git a/pytato/scalar_expr.py b/pytato/scalar_expr.py index 329cd1bd4..b63c1a7c7 100644 --- a/pytato/scalar_expr.py +++ b/pytato/scalar_expr.py @@ -28,7 +28,7 @@ Any, Union, Mapping, FrozenSet, Set, Tuple, Optional, TYPE_CHECKING, Iterable) -import attrs +import dataclasses from pymbolic.mapper import (WalkMapper as WalkMapperBase, IdentityMapper as @@ -263,7 +263,7 @@ def make_stringifier(self, originating_stringifier: Any = None) -> str: return StringifyMapper() -@attrs.frozen(eq=True, hash=True, cache_hash=True) +@dataclasses.dataclass(frozen=True, eq=True, slots=True) class Reduce(ExpressionBase): """ .. autoattribute:: inner_expr @@ -292,14 +292,14 @@ def __getinitargs__(self) -> Tuple[ScalarExpression, ReductionOperation, Any]: return (self.inner_expr, self.op, self.bounds) if __debug__: - def __attrs_post_init__(self) -> None: + def __post_init__(self) -> None: hash(self.bounds) init_arg_names = ("inner_expr", "op", "bounds") mapper_method = "map_reduce" -@attrs.frozen(eq=True, hash=True, cache_hash=True) +@dataclasses.dataclass(frozen=True, eq=True) class TypeCast(ExpressionBase): """ .. autoattribute:: dtype diff --git a/pytato/stringifier.py b/pytato/stringifier.py index f2ffea205..024dc25bc 100644 --- a/pytato/stringifier.py +++ b/pytato/stringifier.py @@ -35,7 +35,7 @@ from pytato.function import FunctionDefinition, Call from pytato.loopy import LoopyCall from immutabledict import immutabledict -import attrs +import dataclasses __doc__ = """ @@ -96,7 +96,7 @@ def _map_generic_array(self, expr: Array, depth: int) -> str: return self.truncation_string # pylint: disable=not-an-iterable - fields = tuple(field.name for field in attrs.fields(type(expr))) + fields = tuple(field.name for field in dataclasses.fields(type(expr))) fields = tuple(field for field in fields if field != "non_equality_tags") @@ -157,7 +157,7 @@ def _get_field_val(field: str) -> str: # pylint: disable=not-an-iterable return (f"{type(expr).__name__}(" + ", ".join(f"{field.name}={_get_field_val(field.name)}" - for field in attrs.fields(type(expr))) + for field in dataclasses.fields(type(expr))) + ")") @memoize_method @@ -174,7 +174,7 @@ def _get_field_val(field: str) -> str: # pylint: disable=not-an-iterable return (f"{type(expr).__name__}(" + ", ".join(f"{field.name}={_get_field_val(field.name)}" - for field in attrs.fields(type(expr))) + for field in dataclasses.fields(type(expr))) + ")") def map_call(self, expr: Call, depth: int) -> str: diff --git a/pytato/target/loopy/codegen.py b/pytato/target/loopy/codegen.py index eb87e44d8..410160a12 100644 --- a/pytato/target/loopy/codegen.py +++ b/pytato/target/loopy/codegen.py @@ -52,7 +52,7 @@ from pytools.tag import Tag import pytato.reductions as red from pytato.codegen import _generate_name_for_temp -import attrs +import dataclasses # set in doc/conf.py if getattr(sys, "_BUILDING_SPHINX_DOCS", False): @@ -107,7 +107,7 @@ def loopy_substitute(expression: Any, variable_assigments: Mapping[str, Any]) -> # {{{ LoopyExpressionContexts -@attrs.define(init=True, repr=False, eq=False) +@dataclasses.dataclass(init=True, repr=False, eq=False) class PersistentExpressionContext(object): """ Mutable state used while generating :mod:`loopy` expressions for a @@ -131,8 +131,7 @@ class PersistentExpressionContext(object): """ state: CodeGenState - _depends_on: FrozenSet[str] = \ - attrs.field(factory=frozenset) + _depends_on: FrozenSet[str] = dataclasses.field(default_factory=frozenset) @property def depends_on(self) -> FrozenSet[str]: @@ -142,7 +141,7 @@ def update_depends_on(self, other: FrozenSet[str]) -> None: self._depends_on = self._depends_on | other -@attrs.define(frozen=True) +@dataclasses.dataclass(frozen=True) class LocalExpressionContext: """ Records context being to be conveyed from a parent expression to its @@ -267,7 +266,7 @@ def to_loopy_expression(self, indices: SymbolicIndex, # {{{ SubstitutionRuleResult -@attrs.define(frozen=True, eq=True) +@dataclasses.dataclass(frozen=True, eq=True) class SubstitutionRuleResult(ImplementedResult): """ An array expression generated as a @@ -291,7 +290,7 @@ def to_loopy_expression(self, # {{{ codegen state -@attrs.define(init=True, repr=False, eq=False) +@dataclasses.dataclass(init=True, repr=False, eq=False) class CodeGenState: """A container for data kept by :class:`CodeGenMapper`. @@ -313,10 +312,10 @@ class CodeGenState: _t_unit: lp.TranslationUnit results: Dict[Array, ImplementedResult] - var_name_gen: pytools.UniqueNameGenerator = attrs.field(init=False) - insn_id_gen: pytools.UniqueNameGenerator = attrs.field(init=False) + var_name_gen: pytools.UniqueNameGenerator = dataclasses.field(init=False) + insn_id_gen: pytools.UniqueNameGenerator = dataclasses.field(init=False) - def __attrs_post_init__(self) -> None: + def __post_init__(self) -> None: self.var_name_gen = self._t_unit.default_entrypoint.get_var_name_generator() self.insn_id_gen = ( self._t_unit.default_entrypoint.get_instruction_id_generator()) @@ -970,7 +969,7 @@ def get_initial_codegen_state(target: LoopyTarget, options=options, lang_version=lp.MOST_RECENT_LANGUAGE_VERSION) - return CodeGenState(t_unit=kernel, results={}) + return CodeGenState(_t_unit=kernel, results={}) # {{{ generate_loopy diff --git a/pytato/transform/__init__.py b/pytato/transform/__init__.py index 76a9ab14b..c9046bf1c 100644 --- a/pytato/transform/__init__.py +++ b/pytato/transform/__init__.py @@ -28,7 +28,7 @@ THE SOFTWARE. """ -import attrs +import dataclasses import logging import numpy as np from immutabledict import immutabledict @@ -382,7 +382,7 @@ def map_loopy_call_result(self, expr: LoopyCallResult) -> Array: rec_container = self.rec(expr._container) assert isinstance(rec_container, LoopyCall) return LoopyCallResult( - container=rec_container, + _container=rec_container, name=expr.name, axes=expr.axes, tags=expr.tags, @@ -422,7 +422,7 @@ def map_function_definition(self, new_mapper = self.clone_for_callee(expr) new_returns = {name: new_mapper(ret) for name, ret in expr.returns.items()} - return attrs.evolve(expr, returns=immutabledict(new_returns)) + return dataclasses.replace(expr, returns=immutabledict(new_returns)) def map_call(self, expr: Call) -> AbstractResultWithNamedArrays: return Call(self.map_function_definition(expr.function), @@ -620,7 +620,7 @@ def map_loopy_call_result(self, expr: LoopyCallResult, rec_loopy_call = self.rec(expr._container, *args, **kwargs) assert isinstance(rec_loopy_call, LoopyCall) return LoopyCallResult( - container=rec_loopy_call, + _container=rec_loopy_call, name=expr.name, axes=expr.axes, tags=expr.tags, diff --git a/pytato/transform/einsum_distributive_law.py b/pytato/transform/einsum_distributive_law.py index 41436c6b6..9df02b36b 100644 --- a/pytato/transform/einsum_distributive_law.py +++ b/pytato/transform/einsum_distributive_law.py @@ -33,7 +33,7 @@ from typing import Callable, Dict, Tuple, Optional, FrozenSet, Mapping -import attrs +import dataclasses from pytato.transform import ArrayOrNames, Mapper, MappedT from pytato.array import (Array, AxesT, Einsum, IndexLambda, EinsumReductionAxis, @@ -54,7 +54,7 @@ class EinsumDistributiveLawDescriptor: """ -@attrs.frozen +@dataclasses.dataclass(frozen=True) class DoNotDistribute(EinsumDistributiveLawDescriptor): """ Tells :func:`apply_distributive_property_to_einsums` to not apply @@ -62,7 +62,7 @@ class DoNotDistribute(EinsumDistributiveLawDescriptor): """ -@attrs.frozen +@dataclasses.dataclass(frozen=True) class DoDistribute(EinsumDistributiveLawDescriptor): """ Tells :func:`apply_distributive_property_to_einsums` to apply distributive @@ -71,17 +71,17 @@ class DoDistribute(EinsumDistributiveLawDescriptor): ioperand: int -@attrs.frozen +@dataclasses.dataclass(frozen=True) class _EinsumDistributiveLawMapperContext: access_descriptors: Tuple[Tuple[EinsumAxisDescriptor, ...], ...] surrounding_args: Mapping[int, Array] redn_axis_to_redn_descr: Mapping[EinsumReductionAxis, ReductionDescriptor] index_to_access_descr: Mapping[str, EinsumAxisDescriptor] - axes: AxesT = attrs.field(kw_only=True) - tags: FrozenSet[Tag] = attrs.field(kw_only=True) + axes: AxesT = dataclasses.field(kw_only=True) + tags: FrozenSet[Tag] = dataclasses.field(kw_only=True) - def __attrs_post_init__(self) -> None: + def __post_init__(self) -> None: # {{{ check that exactly one of the args is missing assert len(self.surrounding_args) == ( diff --git a/pytato/visualization/dot.py b/pytato/visualization/dot.py index 0e3396df0..78cfdc67d 100644 --- a/pytato/visualization/dot.py +++ b/pytato/visualization/dot.py @@ -28,7 +28,7 @@ from functools import partial import html -import attrs +import dataclasses from typing import (TYPE_CHECKING, Callable, Dict, Tuple, Union, List, Mapping, Any, FrozenSet, Set, Optional) @@ -66,7 +66,7 @@ # {{{ _DotEmitter -@attrs.define +@dataclasses.dataclass class _SubgraphTree: contents: Optional[List[str]] subgraphs: Dict[str, _SubgraphTree] @@ -141,7 +141,7 @@ def emit_subgraph(sg: _SubgraphTree) -> None: # {{{ array -> dot node converter -@attrs.define +@dataclasses.dataclass class _DotNodeInfo: title: str fields: Dict[str, Any] @@ -187,7 +187,7 @@ def handle_unsupported_array(self, # type: ignore[override] info = self.get_common_dot_info(expr) # pylint: disable=not-an-iterable - for field in attrs.fields(type(expr)): + for field in dataclasses.fields(type(expr)): if field.name in info.fields: continue attr = getattr(expr, field.name) diff --git a/setup.py b/setup.py index 098aa4a8c..3f27a7b7e 100644 --- a/setup.py +++ b/setup.py @@ -24,20 +24,17 @@ "Programming Language :: Python", "Programming Language :: Python", "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.7", - "Programming Language :: Python :: 3.8", "Topic :: Scientific/Engineering", "Topic :: Scientific/Engineering :: Information Analysis", "Topic :: Scientific/Engineering :: Mathematics", "Topic :: Scientific/Engineering :: Visualization", "Topic :: Software Development :: Libraries", ], - python_requires="~=3.8", + python_requires="~=3.10", install_requires=[ "loopy>=2020.2", "pytools>=2022.1.13", "immutabledict", - "attrs", "bidict", ], package_data={"pytato": ["py.typed"]}, diff --git a/test/test_pytato.py b/test/test_pytato.py index b9fc732b4..1721f3c3b 100644 --- a/test/test_pytato.py +++ b/test/test_pytato.py @@ -29,7 +29,7 @@ import numpy as np import pytest -import attrs +import dataclasses import pytato as pt @@ -1023,14 +1023,14 @@ def test_with_tagged_reduction(): def test_derived_class_uses_correct_array_eq(): - @attrs.define(frozen=True) + @dataclasses.dataclass(frozen=True) class MyNewArrayT(pt.Array): pass with pytest.raises(AssertionError): MyNewArrayT(tags=frozenset(), axes=()) - @attrs.define(frozen=True, eq=False) + @dataclasses.dataclass(frozen=True, eq=False) class MyNewAndCorrectArrayT(pt.Array): pass From 9f7ed0c2963fecea2c44701a379befc79be63080 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Wed, 17 Jul 2024 17:10:37 -0500 Subject: [PATCH 2/9] post-merge fix --- pytato/array.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/pytato/array.py b/pytato/array.py index 38ad18a3b..3337c43f6 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -791,16 +791,16 @@ def __repr__(self) -> str: @dataclasses.dataclass(frozen=True, eq=False, slots=False, repr=False) class _SuppliedAxesAndTagsMixin(Taggable): - axes: AxesT = attrs.field(kw_only=True) - tags: FrozenSet[Tag] = attrs.field(kw_only=True) + axes: AxesT = dataclasses.field(kw_only=True) + tags: FrozenSet[Tag] = dataclasses.field(kw_only=True) # These are automatically excluded from equality in EqualityComparer - non_equality_tags: FrozenSet[Tag] = attrs.field(kw_only=True, + non_equality_tags: FrozenSet[Tag] = dataclasses.field(kw_only=True, hash=False, default=frozenset()) def _with_new_tags(self: Self, tags: FrozenSet[Tag]) -> Self: - return attrs.evolve(self, tags=tags) + return dataclasses.replace(self, tags=tags) @dataclasses.dataclass(frozen=True, eq=False, slots=False, repr=False) From 19286bcc723ff037e2fe056ac980a96ce244588b Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Wed, 17 Jul 2024 17:44:32 -0500 Subject: [PATCH 3/9] more fixes --- pytato/array.py | 2 +- pytato/function.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pytato/array.py b/pytato/array.py index 3337c43f6..24cf0a49f 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -845,7 +845,7 @@ def copy(self, *, # type: ignore[override] if non_equality_tags is None else non_equality_tags) - return type(self)(container=container, + return type(self)(_container=container, name=name, tags=tags, axes=axes, diff --git a/pytato/function.py b/pytato/function.py index 0b8630578..dcbc80512 100644 --- a/pytato/function.py +++ b/pytato/function.py @@ -253,12 +253,12 @@ def without_tags(self, @property def shape(self) -> ShapeType: assert isinstance(self._container, Call) - return self._container.function.returns[self.name].shape + return self._container.function.returns[self.name].shape # pylint: disable=no-member @property def dtype(self) -> _dtype_any: assert isinstance(self._container, Call) - return self._container.function.returns[self.name].dtype + return self._container.function.returns[self.name].dtype # pylint: disable=no-member # eq=False to avoid equality comparison without EqualityMapper From f78cdbffe031aeede2b250ff0171297a4fe831f4 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Mon, 14 Oct 2024 16:34:40 -0500 Subject: [PATCH 4/9] switch to array_dataclass --- pytato/array.py | 49 ++++++++++++++++++++++--------------- pytato/distributed/nodes.py | 7 +++--- pytato/function.py | 8 +++--- 3 files changed, 37 insertions(+), 27 deletions(-) diff --git a/pytato/array.py b/pytato/array.py index f4982c4a2..2b8ee39cc 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -188,6 +188,7 @@ TypeVar, Union, cast, + dataclass_transform, ) import numpy as np @@ -298,6 +299,14 @@ def normalize_shape_component( # {{{ array interface +T = TypeVar("T") + + +@dataclass_transform(eq_default=False, frozen_default=True) +def array_dataclass(cls: type[T]) -> type[T]: + return dataclasses.dataclass(init=True, frozen=True, eq=False, repr=False)(cls) + + ConvertibleToIndexExpr = Union[int, slice, "Array", None, EllipsisType] IndexExpr = Union[IntegralT, "NormalizedSlice", "Array", None, EllipsisType] PyScalarType = Union[type[bool], type[int], type[float], type[complex]] @@ -372,7 +381,7 @@ def _with_new_tags(self, tags: frozenset[Tag]) -> ReductionDescriptor: return replace(self, tags=tags) -@dataclasses.dataclass(frozen=True, eq=False, repr=False, unsafe_hash=True) +@array_dataclass class Array(Taggable): r""" A base class (abstract interface + supplemental functionality) for lazily @@ -735,7 +744,7 @@ def __repr__(self) -> str: # {{{ mixins -@dataclasses.dataclass(frozen=True, eq=False, slots=False, repr=False) +@array_dataclass class _SuppliedAxesAndTagsMixin(Taggable): axes: AxesT = dataclasses.field(kw_only=True) tags: frozenset[Tag] = dataclasses.field(kw_only=True) @@ -749,7 +758,7 @@ def _with_new_tags(self: Self, tags: frozenset[Tag]) -> Self: return dataclasses.replace(self, tags=tags) -@dataclasses.dataclass(frozen=True, eq=False, slots=False, repr=False) +@array_dataclass class _SuppliedShapeAndDtypeMixin: """A mixin class for when an array must store its own *shape* and *dtype*, rather than when it can derive them easily from inputs. @@ -762,7 +771,7 @@ class _SuppliedShapeAndDtypeMixin: # {{{ dict of named arrays -@dataclasses.dataclass(frozen=True, eq=False, repr=False, unsafe_hash=True) +@array_dataclass class NamedArray(_SuppliedAxesAndTagsMixin, Array): """An entry in a :class:`AbstractResultWithNamedArrays`. Holds a reference back to the containing instance as well as the name by which *self* is @@ -813,7 +822,7 @@ def dtype(self) -> np.dtype[Any]: return self.expr.dtype -@dataclasses.dataclass(frozen=True, eq=False, unsafe_hash=True) +@array_dataclass class AbstractResultWithNamedArrays(Mapping[str, NamedArray], Taggable, ABC): r"""An abstract array computation that results in multiple :class:`Array`\ s, each named. The way in which the values of these arrays are computed @@ -863,7 +872,7 @@ def __eq__(self, other: Any) -> bool: return EqualityComparer()(self, other) -@dataclasses.dataclass(frozen=True, eq=False, init=False) +@array_dataclass class DictOfNamedArrays(AbstractResultWithNamedArrays): """A container of named results, each of which can be computed as an array expression provided to the constructor. @@ -918,7 +927,7 @@ def __repr__(self) -> str: # {{{ index lambda -@dataclasses.dataclass(frozen=True, eq=False, repr=False, unsafe_hash=True) +@array_dataclass class IndexLambda(_SuppliedAxesAndTagsMixin, _SuppliedShapeAndDtypeMixin, Array): r"""Represents an array that can be computed by evaluating :attr:`expr` for every value of the input indices. The @@ -1011,7 +1020,7 @@ class EinsumAxisDescriptor: pass -@dataclasses.dataclass(frozen=True, order=True) +@array_dataclass class EinsumElementwiseAxis(EinsumAxisDescriptor): """ Describes an elementwise access pattern of an array's axis. In terms of the @@ -1021,7 +1030,7 @@ class EinsumElementwiseAxis(EinsumAxisDescriptor): dim: int -@dataclasses.dataclass(frozen=True, order=True) +@array_dataclass class EinsumReductionAxis(EinsumAxisDescriptor): """ Describes a reduction access pattern of an array's axis. In terms of the @@ -1031,7 +1040,7 @@ class EinsumReductionAxis(EinsumAxisDescriptor): dim: int -@dataclasses.dataclass(frozen=True, eq=False, repr=False, unsafe_hash=True) +@array_dataclass class Einsum(_SuppliedAxesAndTagsMixin, Array): """ An array expression using the `Einstein summation convention @@ -1368,7 +1377,7 @@ def einsum(subscripts: str, *operands: Array, # {{{ stack -@dataclasses.dataclass(frozen=True, eq=False, repr=False, unsafe_hash=True) +@array_dataclass class Stack(_SuppliedAxesAndTagsMixin, Array): """Join a sequence of arrays along a new axis. @@ -1401,7 +1410,7 @@ def shape(self) -> ShapeType: # {{{ concatenate -@dataclasses.dataclass(frozen=True, eq=False, repr=False, unsafe_hash=True) +@array_dataclass class Concatenate(_SuppliedAxesAndTagsMixin, Array): """Join a sequence of arrays along an existing axis. @@ -1438,7 +1447,7 @@ def shape(self) -> ShapeType: # {{{ index remapping -@dataclasses.dataclass(frozen=True, eq=False, repr=False, unsafe_hash=True) +@array_dataclass class IndexRemappingBase(Array): """Base class for operations that remap the indices of an array. @@ -1461,7 +1470,7 @@ def dtype(self) -> np.dtype[Any]: # {{{ roll -@dataclasses.dataclass(frozen=True, eq=False, repr=False, unsafe_hash=True) +@array_dataclass class Roll(_SuppliedAxesAndTagsMixin, IndexRemappingBase): """Roll an array along an axis. @@ -1487,7 +1496,7 @@ def shape(self) -> ShapeType: # {{{ axis permutation -@dataclasses.dataclass(frozen=True, eq=False, repr=False, unsafe_hash=True) +@array_dataclass class AxisPermutation(_SuppliedAxesAndTagsMixin, IndexRemappingBase): r"""Permute the axes of an array. @@ -1514,7 +1523,7 @@ def shape(self) -> ShapeType: # {{{ reshape -@dataclasses.dataclass(frozen=True, eq=False, repr=False, unsafe_hash=True) +@array_dataclass class Reshape(_SuppliedAxesAndTagsMixin, IndexRemappingBase): """ Reshape an array. @@ -1549,7 +1558,7 @@ def shape(self) -> ShapeType: # {{{ indexing -@dataclasses.dataclass(frozen=True, eq=False, repr=False, unsafe_hash=True) +@array_dataclass class IndexBase(_SuppliedAxesAndTagsMixin, IndexRemappingBase): """ Abstract class for all index expressions on an array. @@ -1669,7 +1678,7 @@ def shape(self) -> ShapeType: # {{{ base class for arguments -@dataclasses.dataclass(frozen=True, eq=False, repr=False, unsafe_hash=True) +@array_dataclass class InputArgumentBase(Array): r"""Base class for input arguments. @@ -1778,7 +1787,7 @@ def dtype(self) -> np.dtype[Any]: # {{{ placeholder -@dataclasses.dataclass(frozen=True, eq=False, repr=False, unsafe_hash=True) +@array_dataclass class Placeholder( _SuppliedAxesAndTagsMixin, _SuppliedShapeAndDtypeMixin, @@ -1802,7 +1811,7 @@ class Placeholder( # {{{ size parameter -@dataclasses.dataclass(frozen=True, eq=False, repr=False, unsafe_hash=True) +@array_dataclass class SizeParam( _SuppliedAxesAndTagsMixin, InputArgumentBase): diff --git a/pytato/distributed/nodes.py b/pytato/distributed/nodes.py index 40ea73cb6..959b7cc9b 100644 --- a/pytato/distributed/nodes.py +++ b/pytato/distributed/nodes.py @@ -70,6 +70,7 @@ _get_default_tags, _SuppliedAxesAndTagsMixin, _SuppliedShapeAndDtypeMixin, + array_dataclass, normalize_shape, ) @@ -79,7 +80,7 @@ # {{{ send -@dataclasses.dataclass(init=True, frozen=True) +@array_dataclass class DistributedSend(Taggable): """Class representing a distributed send operation. See :class:`DistributedSendRefHolder` for a way to ensure that nodes @@ -116,7 +117,7 @@ def copy(self, **kwargs: Any) -> DistributedSend: # {{{ send ref holder -@dataclasses.dataclass(eq=False, frozen=True, repr=False, unsafe_hash=True) +@array_dataclass class DistributedSendRefHolder(Array): """A node acting as an identity on :attr:`passthrough_data` while also holding a reference to a :class:`DistributedSend` in :attr:`send`. Since @@ -181,7 +182,7 @@ def non_equality_tags(self) -> frozenset[Tag]: # {{{ receive -@dataclasses.dataclass(frozen=True, eq=False, unsafe_hash=True) +@array_dataclass class DistributedRecv(_SuppliedAxesAndTagsMixin, _SuppliedShapeAndDtypeMixin, Array): """Class representing a distributed receive operation. diff --git a/pytato/function.py b/pytato/function.py index 2fd0b432e..5ea3ae3b6 100644 --- a/pytato/function.py +++ b/pytato/function.py @@ -84,6 +84,7 @@ Placeholder, ShapeType, _dtype_any, + array_dataclass, ) @@ -103,8 +104,7 @@ class ReturnType(enum.Enum): TUPLE_OF_ARRAYS = 2 -# eq=False to avoid equality comparison without EqualityMapper -@dataclasses.dataclass(frozen=True, eq=False, unsafe_hash=True) +@array_dataclass class FunctionDefinition(Taggable): r""" A function definition that represents its outputs as instances of @@ -246,7 +246,7 @@ def __eq__(self, other: Any) -> bool: return EqualityComparer().map_function_definition(self, other) -@dataclasses.dataclass(frozen=True, eq=False, repr=False, hash=True, cache_hash=True) +@array_dataclass class NamedCallResult(NamedArray): """ One of the arrays that are returned from a call to :class:`FunctionDefinition`. @@ -296,7 +296,7 @@ def dtype(self) -> _dtype_any: # eq=False to avoid equality comparison without EqualityMapper -@dataclasses.dataclass(frozen=True, eq=False, unsafe_hash=True, repr=False) +@array_dataclass class Call(AbstractResultWithNamedArrays): """ Records an invocation to a :class:`FunctionDefinition`. From 459032678e94c2f690083b76d0957029740fa30e Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Wed, 16 Oct 2024 14:20:35 -0500 Subject: [PATCH 5/9] use _augment_array_dataclass --- pytato/array.py | 99 ++++++++++++++++++++++++++---------- pytato/codegen.py | 2 +- pytato/distributed/nodes.py | 8 +-- pytato/function.py | 8 +-- pytato/transform/__init__.py | 2 +- 5 files changed, 82 insertions(+), 37 deletions(-) diff --git a/pytato/array.py b/pytato/array.py index 2b8ee39cc..140aeace9 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -297,15 +297,60 @@ def normalize_shape_component( # }}} -# {{{ array interface +# {{{ array dataclass helpers T = TypeVar("T") @dataclass_transform(eq_default=False, frozen_default=True) -def array_dataclass(cls: type[T]) -> type[T]: - return dataclasses.dataclass(init=True, frozen=True, eq=False, repr=False)(cls) +def array_dataclass() -> Callable[[type[T]], type[T]]: + def map_cls(cls: type[T]) -> type[T]: + # Frozen dataclasses (empirically) have a ~20% speed penalty, + # and their frozen-ness is arguably a debug feature. + dc_cls = dataclasses.dataclass(init=True, frozen=__debug__, + eq=False, repr=False)(cls) + + _augment_array_dataclass(dc_cls) + return dc_cls + + return map_cls + + +def _augment_array_dataclass( + cls: type, + ) -> None: + from dataclasses import fields + attr_tuple = ", ".join(f"self.{fld.name}" + for fld in fields(cls) if fld.name != "non_equality_tags") + if attr_tuple: + attr_tuple = f"({attr_tuple},)" + else: + attr_tuple = "()" + + from pytools.codegen import remove_common_indentation + augment_code = remove_common_indentation( + f""" + def {cls.__name__}_hash(self): + try: + return self._hash_value + except AttributeError: + pass + + h = hash(frozenset({attr_tuple})) + object.__setattr__(self, "_hash_value", h) + return h + + cls.__hash__ = {cls.__name__}_hash + """) + exec_dict = {"cls": cls, "_MODULE_SOURCE_CODE": augment_code} + exec(compile(augment_code, + f"", "exec"), + exec_dict) + +# }}} + +# {{{ array interface ConvertibleToIndexExpr = Union[int, slice, "Array", None, EllipsisType] IndexExpr = Union[IntegralT, "NormalizedSlice", "Array", None, EllipsisType] @@ -381,7 +426,7 @@ def _with_new_tags(self, tags: frozenset[Tag]) -> ReductionDescriptor: return replace(self, tags=tags) -@array_dataclass +@array_dataclass() class Array(Taggable): r""" A base class (abstract interface + supplemental functionality) for lazily @@ -555,7 +600,7 @@ def ndim(self) -> int: @property def T(self) -> Array: return AxisPermutation(self, - tuple(range(self.ndim)[::-1]), + axis_permutation=tuple(range(self.ndim)[::-1]), tags=_get_default_tags(), axes=_get_default_axes(self.ndim), non_equality_tags=_get_created_at_tag()) @@ -744,7 +789,7 @@ def __repr__(self) -> str: # {{{ mixins -@array_dataclass +@dataclasses.dataclass(frozen=True, eq=False, repr=False) class _SuppliedAxesAndTagsMixin(Taggable): axes: AxesT = dataclasses.field(kw_only=True) tags: frozenset[Tag] = dataclasses.field(kw_only=True) @@ -758,7 +803,7 @@ def _with_new_tags(self: Self, tags: frozenset[Tag]) -> Self: return dataclasses.replace(self, tags=tags) -@array_dataclass +@dataclasses.dataclass(frozen=True, eq=False, repr=False) class _SuppliedShapeAndDtypeMixin: """A mixin class for when an array must store its own *shape* and *dtype*, rather than when it can derive them easily from inputs. @@ -771,7 +816,7 @@ class _SuppliedShapeAndDtypeMixin: # {{{ dict of named arrays -@array_dataclass +@array_dataclass() class NamedArray(_SuppliedAxesAndTagsMixin, Array): """An entry in a :class:`AbstractResultWithNamedArrays`. Holds a reference back to the containing instance as well as the name by which *self* is @@ -822,7 +867,7 @@ def dtype(self) -> np.dtype[Any]: return self.expr.dtype -@array_dataclass +@dataclasses.dataclass(frozen=True, eq=False) class AbstractResultWithNamedArrays(Mapping[str, NamedArray], Taggable, ABC): r"""An abstract array computation that results in multiple :class:`Array`\ s, each named. The way in which the values of these arrays are computed @@ -872,7 +917,7 @@ def __eq__(self, other: Any) -> bool: return EqualityComparer()(self, other) -@array_dataclass +@dataclasses.dataclass(frozen=True, eq=False, init=False) class DictOfNamedArrays(AbstractResultWithNamedArrays): """A container of named results, each of which can be computed as an array expression provided to the constructor. @@ -927,7 +972,7 @@ def __repr__(self) -> str: # {{{ index lambda -@array_dataclass +@array_dataclass() class IndexLambda(_SuppliedAxesAndTagsMixin, _SuppliedShapeAndDtypeMixin, Array): r"""Represents an array that can be computed by evaluating :attr:`expr` for every value of the input indices. The @@ -1020,7 +1065,7 @@ class EinsumAxisDescriptor: pass -@array_dataclass +@dataclasses.dataclass(frozen=True, order=True) class EinsumElementwiseAxis(EinsumAxisDescriptor): """ Describes an elementwise access pattern of an array's axis. In terms of the @@ -1030,7 +1075,7 @@ class EinsumElementwiseAxis(EinsumAxisDescriptor): dim: int -@array_dataclass +@dataclasses.dataclass(frozen=True, order=True) class EinsumReductionAxis(EinsumAxisDescriptor): """ Describes a reduction access pattern of an array's axis. In terms of the @@ -1040,7 +1085,7 @@ class EinsumReductionAxis(EinsumAxisDescriptor): dim: int -@array_dataclass +@array_dataclass() class Einsum(_SuppliedAxesAndTagsMixin, Array): """ An array expression using the `Einstein summation convention @@ -1377,7 +1422,7 @@ def einsum(subscripts: str, *operands: Array, # {{{ stack -@array_dataclass +@array_dataclass() class Stack(_SuppliedAxesAndTagsMixin, Array): """Join a sequence of arrays along a new axis. @@ -1410,7 +1455,7 @@ def shape(self) -> ShapeType: # {{{ concatenate -@array_dataclass +@array_dataclass() class Concatenate(_SuppliedAxesAndTagsMixin, Array): """Join a sequence of arrays along an existing axis. @@ -1447,7 +1492,7 @@ def shape(self) -> ShapeType: # {{{ index remapping -@array_dataclass +@array_dataclass() class IndexRemappingBase(Array): """Base class for operations that remap the indices of an array. @@ -1470,7 +1515,7 @@ def dtype(self) -> np.dtype[Any]: # {{{ roll -@array_dataclass +@array_dataclass() class Roll(_SuppliedAxesAndTagsMixin, IndexRemappingBase): """Roll an array along an axis. @@ -1496,7 +1541,7 @@ def shape(self) -> ShapeType: # {{{ axis permutation -@array_dataclass +@array_dataclass() class AxisPermutation(_SuppliedAxesAndTagsMixin, IndexRemappingBase): r"""Permute the axes of an array. @@ -1523,7 +1568,7 @@ def shape(self) -> ShapeType: # {{{ reshape -@array_dataclass +@array_dataclass() class Reshape(_SuppliedAxesAndTagsMixin, IndexRemappingBase): """ Reshape an array. @@ -1558,7 +1603,7 @@ def shape(self) -> ShapeType: # {{{ indexing -@array_dataclass +@array_dataclass() class IndexBase(_SuppliedAxesAndTagsMixin, IndexRemappingBase): """ Abstract class for all index expressions on an array. @@ -1678,7 +1723,7 @@ def shape(self) -> ShapeType: # {{{ base class for arguments -@array_dataclass +@array_dataclass() class InputArgumentBase(Array): r"""Base class for input arguments. @@ -1787,7 +1832,7 @@ def dtype(self) -> np.dtype[Any]: # {{{ placeholder -@array_dataclass +@array_dataclass() class Placeholder( _SuppliedAxesAndTagsMixin, _SuppliedShapeAndDtypeMixin, @@ -1811,7 +1856,7 @@ class Placeholder( # {{{ size parameter -@array_dataclass +@array_dataclass() class SizeParam( _SuppliedAxesAndTagsMixin, InputArgumentBase): @@ -1823,8 +1868,8 @@ class SizeParam( The name by which a value is supplied for the argument once computation begins. """ - name: str - axes: AxesT = dataclasses.field(kw_only=True, default=()) + name: str = dataclasses.field(kw_only=True) # pylint: disable=invalid-field-call + axes: AxesT = dataclasses.field(kw_only=True, default=()) # pylint: disable=invalid-field-call _mapper_method: ClassVar[str] = "map_size_param" @@ -2169,7 +2214,7 @@ def make_size_param(name: str, :param tags: implementation tags """ _check_identifier(name, optional=False) - return SizeParam(name, tags=(tags | _get_default_tags()), + return SizeParam(name=name, tags=(tags | _get_default_tags()), # pylint: disable=missing-kwoa non_equality_tags=_get_created_at_tag(),) diff --git a/pytato/codegen.py b/pytato/codegen.py index 48ed49cb2..019b45e52 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -130,7 +130,7 @@ def __init__(self, target: Target, def map_size_param(self, expr: SizeParam) -> Array: name = expr.name assert name is not None - return SizeParam( + return SizeParam( # pylint: disable=missing-kwoa name=name, tags=expr.tags, non_equality_tags=expr.non_equality_tags) diff --git a/pytato/distributed/nodes.py b/pytato/distributed/nodes.py index 959b7cc9b..7daa5284d 100644 --- a/pytato/distributed/nodes.py +++ b/pytato/distributed/nodes.py @@ -80,7 +80,7 @@ # {{{ send -@array_dataclass +@array_dataclass() class DistributedSend(Taggable): """Class representing a distributed send operation. See :class:`DistributedSendRefHolder` for a way to ensure that nodes @@ -104,7 +104,7 @@ class DistributedSend(Taggable): data: Array dest_rank: int comm_tag: CommTagType - tags: frozenset[Tag] = dataclasses.field(kw_only=True, default=frozenset()) + tags: frozenset[Tag] = dataclasses.field(kw_only=True, default=frozenset()) # pylint: disable=invalid-field-call def _with_new_tags(self, tags: frozenset[Tag]) -> DistributedSend: return dataclasses.replace(self, tags=tags) @@ -117,7 +117,7 @@ def copy(self, **kwargs: Any) -> DistributedSend: # {{{ send ref holder -@array_dataclass +@array_dataclass() class DistributedSendRefHolder(Array): """A node acting as an identity on :attr:`passthrough_data` while also holding a reference to a :class:`DistributedSend` in :attr:`send`. Since @@ -182,7 +182,7 @@ def non_equality_tags(self) -> frozenset[Tag]: # {{{ receive -@array_dataclass +@array_dataclass() class DistributedRecv(_SuppliedAxesAndTagsMixin, _SuppliedShapeAndDtypeMixin, Array): """Class representing a distributed receive operation. diff --git a/pytato/function.py b/pytato/function.py index 5ea3ae3b6..b27fd9c9c 100644 --- a/pytato/function.py +++ b/pytato/function.py @@ -104,7 +104,7 @@ class ReturnType(enum.Enum): TUPLE_OF_ARRAYS = 2 -@array_dataclass +@array_dataclass() class FunctionDefinition(Taggable): r""" A function definition that represents its outputs as instances of @@ -158,7 +158,7 @@ class FunctionDefinition(Taggable): parameters: frozenset[str] return_type: ReturnType returns: Mapping[str, Array] - tags: frozenset[Tag] = dataclasses.field(kw_only=True) + tags: frozenset[Tag] = dataclasses.field(kw_only=True) # pylint: disable=invalid-field-call @cached_property def _placeholders(self) -> Mapping[str, Placeholder]: @@ -246,7 +246,7 @@ def __eq__(self, other: Any) -> bool: return EqualityComparer().map_function_definition(self, other) -@array_dataclass +@array_dataclass() class NamedCallResult(NamedArray): """ One of the arrays that are returned from a call to :class:`FunctionDefinition`. @@ -296,7 +296,7 @@ def dtype(self) -> _dtype_any: # eq=False to avoid equality comparison without EqualityMapper -@array_dataclass +@array_dataclass() class Call(AbstractResultWithNamedArrays): """ Records an invocation to a :class:`FunctionDefinition`. diff --git a/pytato/transform/__init__.py b/pytato/transform/__init__.py index af680c1fd..584f6bf86 100644 --- a/pytato/transform/__init__.py +++ b/pytato/transform/__init__.py @@ -646,7 +646,7 @@ def map_data_wrapper(self, expr: DataWrapper, def map_size_param(self, expr: SizeParam, *args: Any, **kwargs: Any) -> Array: assert expr.name is not None - return SizeParam(expr.name, axes=expr.axes, tags=expr.tags) + return SizeParam(name=expr.name, axes=expr.axes, tags=expr.tags) def map_einsum(self, expr: Einsum, *args: Any, **kwargs: Any) -> Array: return Einsum(expr.access_descriptors, From 2da059e057927f6e89f61c4d5f025e871a0625ca Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Wed, 16 Oct 2024 14:55:21 -0500 Subject: [PATCH 6/9] restore attribute validation --- pytato/array.py | 11 +++++++++++ pytato/function.py | 6 +++++- pytato/loopy.py | 4 ++++ 3 files changed, 20 insertions(+), 1 deletion(-) diff --git a/pytato/array.py b/pytato/array.py index 140aeace9..0699e3228 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -1013,6 +1013,12 @@ class IndexLambda(_SuppliedAxesAndTagsMixin, _SuppliedShapeAndDtypeMixin, Array) _mapper_method: ClassVar[str] = "map_index_lambda" + if __debug__: + def __post_init__(self) -> None: + assert isinstance(self.bindings, immutabledict) + assert isinstance(self.var_to_reduction_descr, immutabledict) + super().__post_init__() + def with_tagged_reduction(self, reduction_variable: str, tag: Tag) -> IndexLambda: @@ -1125,6 +1131,11 @@ class Einsum(_SuppliedAxesAndTagsMixin, Array): ReductionDescriptor] _mapper_method: ClassVar[str] = "map_einsum" + if __debug__: + def __post_init__(self) -> None: + assert isinstance(self.redn_axis_to_redn_descr, immutabledict) + super().__post_init__() + @memoize_method def _access_descr_to_axis_len(self ) -> Mapping[EinsumAxisDescriptor, ShapeComponent]: diff --git a/pytato/function.py b/pytato/function.py index b27fd9c9c..f04b3b926 100644 --- a/pytato/function.py +++ b/pytato/function.py @@ -160,6 +160,10 @@ class FunctionDefinition(Taggable): returns: Mapping[str, Array] tags: frozenset[Tag] = dataclasses.field(kw_only=True) # pylint: disable=invalid-field-call + if __debug__: + def __post_init__(self) -> None: + assert isinstance(self.returns, immutabledict) + @cached_property def _placeholders(self) -> Mapping[str, Placeholder]: from pytato.transform import InputGatherer @@ -295,7 +299,6 @@ def dtype(self) -> _dtype_any: return self._container.function.returns[self.name].dtype # pylint: disable=no-member -# eq=False to avoid equality comparison without EqualityMapper @array_dataclass() class Call(AbstractResultWithNamedArrays): """ @@ -323,6 +326,7 @@ def __post_init__(self) -> None: # check that the invocation parameters and the function definition # parameters agree with each other. assert frozenset(self.bindings) == self.function.parameters + assert isinstance(self.bindings, immutabledict) super().__post_init__() def __contains__(self, name: object) -> bool: diff --git a/pytato/loopy.py b/pytato/loopy.py index 92c6527e9..a931250ad 100644 --- a/pytato/loopy.py +++ b/pytato/loopy.py @@ -106,6 +106,10 @@ class LoopyCall(AbstractResultWithNamedArrays): copy = dataclasses.replace + def __post_init__(self) -> None: + assert isinstance(self.bindings, immutabledict) + super().__post_init__() + @property def _result_names(self) -> frozenset[str]: return frozenset({name From 1a646954909cb9ce89ddecc68cad98a5db05bee6 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Wed, 13 Nov 2024 15:23:41 -0600 Subject: [PATCH 7/9] mypy? --- pytato/loopy.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pytato/loopy.py b/pytato/loopy.py index acb6f4d5a..f57f7d6d5 100644 --- a/pytato/loopy.py +++ b/pytato/loopy.py @@ -52,6 +52,7 @@ NamedArray, ShapeType, SizeParam, + array_dataclass, ) from pytato.scalar_expr import ( EvaluationMapper, @@ -91,7 +92,7 @@ """ -@dataclasses.dataclass(frozen=True, eq=False) +@array_dataclass() class LoopyCall(AbstractResultWithNamedArrays): """ An array expression node representing a call to an entrypoint in a From c8a14b76c395b7fc414607deb705e6ca98764a9f Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Tue, 19 Nov 2024 16:23:09 -0600 Subject: [PATCH 8/9] Disable pylint invalid-field-call --- .pylintrc-local.yml | 1 + pytato/array.py | 4 ++-- pytato/distributed/nodes.py | 2 +- pytato/function.py | 2 +- 4 files changed, 5 insertions(+), 4 deletions(-) diff --git a/.pylintrc-local.yml b/.pylintrc-local.yml index ae0f92032..b0c3451ab 100644 --- a/.pylintrc-local.yml +++ b/.pylintrc-local.yml @@ -13,3 +13,4 @@ - too-many-function-args - redundant-keyword-arg - no-value-for-parameter + - invalid-field-call diff --git a/pytato/array.py b/pytato/array.py index 8dbc06268..6cc998812 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -1886,8 +1886,8 @@ class SizeParam( The name by which a value is supplied for the argument once computation begins. """ - name: str = dataclasses.field(kw_only=True) # pylint: disable=invalid-field-call - axes: AxesT = dataclasses.field(kw_only=True, default=()) # pylint: disable=invalid-field-call + name: str = dataclasses.field(kw_only=True) + axes: AxesT = dataclasses.field(kw_only=True, default=()) _mapper_method: ClassVar[str] = "map_size_param" diff --git a/pytato/distributed/nodes.py b/pytato/distributed/nodes.py index 95e032e4f..de06f7578 100644 --- a/pytato/distributed/nodes.py +++ b/pytato/distributed/nodes.py @@ -105,7 +105,7 @@ class DistributedSend(Taggable): data: Array dest_rank: int comm_tag: CommTagType - tags: frozenset[Tag] = dataclasses.field(kw_only=True, default=frozenset()) # pylint: disable=invalid-field-call + tags: frozenset[Tag] = dataclasses.field(kw_only=True, default=frozenset()) def _with_new_tags(self, tags: frozenset[Tag]) -> DistributedSend: return dataclasses.replace(self, tags=tags) diff --git a/pytato/function.py b/pytato/function.py index 6554d660a..96fef66b2 100644 --- a/pytato/function.py +++ b/pytato/function.py @@ -152,7 +152,7 @@ class FunctionDefinition(Taggable): parameters: frozenset[str] return_type: ReturnType returns: Mapping[str, Array] - tags: frozenset[Tag] = dataclasses.field(kw_only=True) # pylint: disable=invalid-field-call + tags: frozenset[Tag] = dataclasses.field(kw_only=True) if __debug__: def __post_init__(self) -> None: From f364e9ce068751d32732d92aa3253b5cb05e08e9 Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Tue, 19 Nov 2024 16:32:22 -0600 Subject: [PATCH 9/9] Automatically generate _mapper_method --- pytato/array.py | 108 +++++++++++++++++++++--------------- pytato/distributed/nodes.py | 6 +- pytato/function.py | 4 -- pytato/loopy.py | 10 ++-- 4 files changed, 69 insertions(+), 59 deletions(-) diff --git a/pytato/array.py b/pytato/array.py index 6cc998812..b637197d9 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -179,6 +179,7 @@ from abc import ABC, abstractmethod from collections.abc import Callable, Collection, Iterable, Iterator, Mapping, Sequence from functools import cached_property, partialmethod +from sys import intern from typing import ( TYPE_CHECKING, Any, @@ -190,6 +191,7 @@ cast, dataclass_transform, ) +from warnings import warn import numpy as np from immutabledict import immutabledict @@ -300,21 +302,39 @@ def normalize_shape_component( @dataclass_transform(eq_default=False, frozen_default=True) -def array_dataclass() -> Callable[[type[T]], type[T]]: +def array_dataclass(hash: bool = True) -> Callable[[type[T]], type[T]]: def map_cls(cls: type[T]) -> type[T]: # Frozen dataclasses (empirically) have a ~20% speed penalty, # and their frozen-ness is arguably a debug feature. dc_cls = dataclasses.dataclass(init=True, frozen=__debug__, eq=False, repr=False)(cls) - _augment_array_dataclass(dc_cls) + _augment_array_dataclass(dc_cls, generate_hash=hash) return dc_cls return map_cls +# https://stackoverflow.com/a/1176023 +_CAMEL_TO_SNAKE_RE = re.compile( + r""" + (?<=[a-z]) # preceded by lowercase + (?=[A-Z]) # followed by uppercase + | # OR + (?<=[A-Z]) # preceded by lowercase + (?=[A-Z][a-z]) # followed by uppercase, then lowercase + """, + re.X, +) + + +class _HasMapperMethod(Protocol): + _mapper_method: ClassVar[str] + + def _augment_array_dataclass( cls: type, + generate_hash: bool, ) -> None: from dataclasses import fields attr_tuple = ", ".join(f"self.{fld.name}" @@ -324,25 +344,48 @@ def _augment_array_dataclass( else: attr_tuple = "()" - from pytools.codegen import remove_common_indentation - augment_code = remove_common_indentation( - f""" - def {cls.__name__}_hash(self): - try: - return self._hash_value - except AttributeError: - pass + if generate_hash: + from pytools.codegen import remove_common_indentation + augment_code = remove_common_indentation( + f""" + def {cls.__name__}_hash(self): + try: + return self._hash_value + except AttributeError: + pass + + h = hash(frozenset({attr_tuple})) + object.__setattr__(self, "_hash_value", h) + return h + + cls.__hash__ = {cls.__name__}_hash + """) + exec_dict = {"cls": cls, "_MODULE_SOURCE_CODE": augment_code} + exec(compile(augment_code, + f"", "exec"), + exec_dict) + + # {{{ assign mapper_method + + mm_cls = cast(type[_HasMapperMethod], cls) + + snake_clsname = _CAMEL_TO_SNAKE_RE.sub("_", mm_cls.__name__).lower() + default_mapper_method_name = f"map_{snake_clsname}" + + # This covers two cases: the class does not have the attribute in the first + # place, or it inherits a value but does not set it itself. + sets_mapper_method = "_mapper_method" in mm_cls.__dict__ - h = hash(frozenset({attr_tuple})) - object.__setattr__(self, "_hash_value", h) - return h + if sets_mapper_method: + if default_mapper_method_name == mm_cls._mapper_method: + warn(f"Explicit _mapper_method on {mm_cls} not needed, default matches " + "explicit assignment. Just delete the explicit assignment.", + stacklevel=3) - cls.__hash__ = {cls.__name__}_hash - """) - exec_dict = {"cls": cls, "_MODULE_SOURCE_CODE": augment_code} - exec(compile(augment_code, - f"", "exec"), - exec_dict) + if not sets_mapper_method: + mm_cls._mapper_method = intern(default_mapper_method_name) + + # }}} # }}} @@ -831,8 +874,6 @@ class NamedArray(_SuppliedAxesAndTagsMixin, Array): _container: AbstractResultWithNamedArrays name: str - _mapper_method: ClassVar[str] = "map_named_array" - # type-ignore reason: `copy` signature incompatible with super-class def copy(self, *, # type: ignore[override] container: AbstractResultWithNamedArrays | None = None, @@ -937,7 +978,6 @@ class DictOfNamedArrays(AbstractResultWithNamedArrays): def __init__(self, data: Mapping[str, Array], *, tags: frozenset[Tag] | None = None) -> None: if tags is None: - from warnings import warn warn("Passing `tags=None` is deprecated and will result" " in an error from 2023. To remove this message either" " call make_dict_of_named_arrays or pass the `tags` argument.", @@ -1015,8 +1055,6 @@ class IndexLambda(_SuppliedAxesAndTagsMixin, _SuppliedShapeAndDtypeMixin, Array) bindings: Mapping[str, Array] var_to_reduction_descr: Mapping[str, ReductionDescriptor] - _mapper_method: ClassVar[str] = "map_index_lambda" - if __debug__: def __post_init__(self) -> None: assert isinstance(self.bindings, immutabledict) @@ -1133,7 +1171,6 @@ class Einsum(_SuppliedAxesAndTagsMixin, Array): args: tuple[Array, ...] redn_axis_to_redn_descr: Mapping[EinsumReductionAxis, ReductionDescriptor] - _mapper_method: ClassVar[str] = "map_einsum" if __debug__: def __post_init__(self) -> None: @@ -1453,8 +1490,6 @@ class Stack(_SuppliedAxesAndTagsMixin, Array): arrays: tuple[Array, ...] axis: int - _mapper_method: ClassVar[str] = "map_stack" - @property def dtype(self) -> np.dtype[Any]: return _np_result_dtype(*(arr.dtype for arr in self.arrays)) @@ -1486,8 +1521,6 @@ class Concatenate(_SuppliedAxesAndTagsMixin, Array): arrays: tuple[Array, ...] axis: int - _mapper_method: ClassVar[str] = "map_concatenate" - @property def dtype(self) -> np.dtype[Any]: return _np_result_dtype(*(arr.dtype for arr in self.arrays)) @@ -1545,8 +1578,6 @@ class Roll(_SuppliedAxesAndTagsMixin, IndexRemappingBase): shift: int axis: int - _mapper_method: ClassVar[str] = "map_roll" - @property def shape(self) -> ShapeType: return self.array.shape @@ -1568,8 +1599,6 @@ class AxisPermutation(_SuppliedAxesAndTagsMixin, IndexRemappingBase): """ axis_permutation: tuple[int, ...] - _mapper_method: ClassVar[str] = "map_axis_permutation" - @property def shape(self) -> ShapeType: result = [] @@ -1603,8 +1632,6 @@ class Reshape(_SuppliedAxesAndTagsMixin, IndexRemappingBase): newshape: ShapeType order: str - _mapper_method: ClassVar[str] = "map_reshape" - if __debug__: def __post_init__(self) -> None: super().__post_init__() @@ -1628,12 +1655,12 @@ class IndexBase(_SuppliedAxesAndTagsMixin, IndexRemappingBase): indices: tuple[IndexExpr, ...] +@array_dataclass() class BasicIndex(IndexBase): """ An indexing expression with all indices being either an :class:`int` or :class:`slice`. """ - _mapper_method: ClassVar[str] = "map_basic_index" @property def shape(self) -> ShapeType: @@ -1781,7 +1808,7 @@ def dtype(self) -> np.dtype[Any]: pass -@dataclasses.dataclass(frozen=True, eq=False, repr=False) +@array_dataclass(hash=False) class DataWrapper(_SuppliedAxesAndTagsMixin, InputArgumentBase): """Takes concrete array data and packages it to be compatible with the :class:`Array` interface. @@ -1824,8 +1851,6 @@ class DataWrapper(_SuppliedAxesAndTagsMixin, InputArgumentBase): data: DataInterface shape: ShapeType - _mapper_method: ClassVar[str] = "map_data_wrapper" - @property def name(self) -> None: return None @@ -1867,8 +1892,6 @@ class Placeholder( """ name: str - _mapper_method: ClassVar[str] = "map_placeholder" - # }}} @@ -1889,8 +1912,6 @@ class SizeParam( name: str = dataclasses.field(kw_only=True) axes: AxesT = dataclasses.field(kw_only=True, default=()) - _mapper_method: ClassVar[str] = "map_size_param" - @property def shape(self) -> ShapeType: return () @@ -2254,7 +2275,6 @@ def make_data_wrapper(data: DataInterface, shape = data.shape if name is not None: - from warnings import warn warn("Naming DataWrappers is deprecated and " "will be converted to a PrefixNamed tag. " "This will stop working in 2023. " diff --git a/pytato/distributed/nodes.py b/pytato/distributed/nodes.py index de06f7578..c7edf7616 100644 --- a/pytato/distributed/nodes.py +++ b/pytato/distributed/nodes.py @@ -55,7 +55,7 @@ import dataclasses from collections.abc import Hashable -from typing import Any, ClassVar +from typing import Any import numpy as np @@ -156,8 +156,6 @@ class DistributedSendRefHolder(Array): send: DistributedSend passthrough_data: Array - _mapper_method: ClassVar[str] = "map_distributed_send_ref_holder" - @property def shape(self) -> ShapeType: return self.passthrough_data.shape @@ -213,8 +211,6 @@ class DistributedRecv(_SuppliedAxesAndTagsMixin, _SuppliedShapeAndDtypeMixin, Ar src_rank: int comm_tag: CommTagType - _mapper_method: ClassVar[str] = "map_distributed_recv" - # }}} diff --git a/pytato/function.py b/pytato/function.py index 96fef66b2..3450e62e6 100644 --- a/pytato/function.py +++ b/pytato/function.py @@ -62,7 +62,6 @@ from functools import cached_property from typing import ( Any, - ClassVar, TypeVar, ) @@ -258,7 +257,6 @@ class NamedCallResult(NamedArray): The name by which the returned array is referred to in :attr:`FunctionDefinition.returns`. """ - _mapper_method: ClassVar[str] = "map_named_call_result" def with_tagged_axis(self, iaxis: int, tags: Iterable[Tag] | Tag) -> Array: @@ -311,8 +309,6 @@ class Call(AbstractResultWithNamedArrays): function: FunctionDefinition bindings: Mapping[str, Array] - _mapper_method: ClassVar[str] = "map_call" - copy = dataclasses.replace if __debug__: diff --git a/pytato/loopy.py b/pytato/loopy.py index f57f7d6d5..118979eea 100644 --- a/pytato/loopy.py +++ b/pytato/loopy.py @@ -33,7 +33,6 @@ from numbers import Number from typing import ( Any, - ClassVar, ) import islpy as isl @@ -102,8 +101,6 @@ class LoopyCall(AbstractResultWithNamedArrays): bindings: Mapping[str, ArrayOrScalar] entrypoint: str - _mapper_method: ClassVar[str] = "map_loopy_call" - copy = dataclasses.replace def __post_init__(self) -> None: @@ -155,13 +152,14 @@ def __iter__(self) -> Iterator[str]: return iter(self._result_names) -@dataclasses.dataclass(frozen=True, eq=False) -class LoopyCallResult(NamedArray): +@array_dataclass() +# https://github.com/python/mypy/issues/18115 +# https://github.com/python/mypy/issues/17623 +class LoopyCallResult(NamedArray): # type: ignore[override] """ Named array for :class:`LoopyCall`'s result. Inherits from :class:`~pytato.array.NamedArray`. """ - _mapper_method = "map_loopy_call_result" _container: LoopyCall @property