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/doc/conf.py b/doc/conf.py index cbc96fbbb..1c69ebfa6 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -28,7 +28,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/pyproject.toml b/pyproject.toml index 9f5ecb578..a029d972f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,7 +32,6 @@ classifiers = [ "Topic :: Software Development :: Libraries", ] dependencies = [ - "attrs", "bidict", "immutabledict", "loopy>=2020.2", diff --git a/pytato/array.py b/pytato/array.py index fc392ac61..b637197d9 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -173,11 +173,13 @@ # }}} +import dataclasses import operator import re 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, @@ -187,9 +189,10 @@ TypeVar, Union, cast, + dataclass_transform, ) +from warnings import warn -import attrs import numpy as np from immutabledict import immutabledict from typing_extensions import Self @@ -293,6 +296,100 @@ def normalize_shape_component( # }}} +# {{{ array dataclass helpers + +T = TypeVar("T") + + +@dataclass_transform(eq_default=False, frozen_default=True) +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, 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}" + for fld in fields(cls) if fld.name != "non_equality_tags") + if attr_tuple: + attr_tuple = f"({attr_tuple},)" + else: + attr_tuple = "()" + + 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__ + + 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) + + if not sets_mapper_method: + mm_cls._mapper_method = intern(default_mapper_method_name) + + # }}} + +# }}} + + # {{{ array interface ConvertibleToIndexExpr = Union[int, slice, "Array", EllipsisType, None] @@ -317,7 +414,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 @@ -342,7 +439,7 @@ class NormalizedSlice: step: IntegerT -@attrs.frozen +@dataclasses.dataclass(frozen=True) class Axis(Taggable): """ A type for recording the information about an :class:`~pytato.Array`'s @@ -351,11 +448,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 @@ -364,11 +461,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) +@array_dataclass() class Array(Taggable): r""" A base class (abstract interface + supplemental functionality) for lazily @@ -458,7 +555,6 @@ class Array(Taggable): .. attribute:: ndim """ - # otherwise subclasses cannot set these if TYPE_CHECKING: @property @@ -487,7 +583,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( @@ -503,7 +599,7 @@ 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) @property def size(self) -> ShapeComponent: @@ -543,7 +639,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()) @@ -740,21 +836,21 @@ def __repr__(self) -> str: # {{{ mixins -@attrs.frozen(eq=False, slots=False, repr=False) +@dataclasses.dataclass(frozen=True, eq=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) -@attrs.frozen(eq=False, slots=False, repr=False) +@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. @@ -767,7 +863,7 @@ class _SuppliedShapeAndDtypeMixin: # {{{ dict of named arrays -@attrs.frozen(eq=False, repr=False, hash=True, cache_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 @@ -778,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, @@ -796,7 +890,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, @@ -818,7 +912,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) 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 @@ -834,14 +928,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() @@ -868,7 +962,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. @@ -877,15 +971,13 @@ 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" 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.", @@ -924,7 +1016,7 @@ def __repr__(self) -> str: # {{{ index lambda -@attrs.frozen(eq=False, repr=False, hash=True, cache_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 @@ -960,12 +1052,14 @@ class IndexLambda(_SuppliedAxesAndTagsMixin, _SuppliedShapeAndDtypeMixin, Array) .. automethod:: with_tagged_reduction """ expr: ScalarExpression - 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" + 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, @@ -1019,7 +1113,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 @@ -1029,7 +1123,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 @@ -1039,7 +1133,7 @@ class EinsumReductionAxis(EinsumAxisDescriptor): dim: int -@attrs.frozen(frozen=True, eq=False, repr=False, hash=True, cache_hash=True) +@array_dataclass() class Einsum(_SuppliedAxesAndTagsMixin, Array): """ An array expression using the `Einstein summation convention @@ -1076,9 +1170,12 @@ class Einsum(_SuppliedAxesAndTagsMixin, 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)) - _mapper_method: ClassVar[str] = "map_einsum" + ReductionDescriptor] + + 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 @@ -1377,7 +1474,7 @@ def einsum(subscripts: str, *operands: Array, # {{{ stack -@attrs.frozen(eq=False, repr=False, hash=True, cache_hash=True) +@array_dataclass() class Stack(_SuppliedAxesAndTagsMixin, Array): """Join a sequence of arrays along a new axis. @@ -1393,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)) @@ -1410,7 +1505,7 @@ def shape(self) -> ShapeType: # {{{ concatenate -@attrs.frozen(eq=False, repr=False, hash=True, cache_hash=True) +@array_dataclass() class Concatenate(_SuppliedAxesAndTagsMixin, Array): """Join a sequence of arrays along an existing axis. @@ -1426,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)) @@ -1447,7 +1540,7 @@ def shape(self) -> ShapeType: # {{{ index remapping -@attrs.frozen(eq=False, repr=False, hash=True, cache_hash=True) +@array_dataclass() class IndexRemappingBase(Array): """Base class for operations that remap the indices of an array. @@ -1470,7 +1563,7 @@ def dtype(self) -> np.dtype[Any]: # {{{ roll -@attrs.frozen(eq=False, repr=False, hash=True, cache_hash=True) +@array_dataclass() class Roll(_SuppliedAxesAndTagsMixin, IndexRemappingBase): """Roll an array along an axis. @@ -1485,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 @@ -1496,7 +1587,7 @@ def shape(self) -> ShapeType: # {{{ axis permutation -@attrs.frozen(eq=False, repr=False, hash=True, cache_hash=True) +@array_dataclass() class AxisPermutation(_SuppliedAxesAndTagsMixin, IndexRemappingBase): r"""Permute the axes of an array. @@ -1508,8 +1599,6 @@ class AxisPermutation(_SuppliedAxesAndTagsMixin, IndexRemappingBase): """ axis_permutation: tuple[int, ...] - _mapper_method: ClassVar[str] = "map_axis_permutation" - @property def shape(self) -> ShapeType: result = [] @@ -1523,7 +1612,7 @@ def shape(self) -> ShapeType: # {{{ reshape -@attrs.frozen(eq=False, repr=False, hash=True, cache_hash=True) +@array_dataclass() class Reshape(_SuppliedAxesAndTagsMixin, IndexRemappingBase): """ Reshape an array. @@ -1543,11 +1632,9 @@ class Reshape(_SuppliedAxesAndTagsMixin, IndexRemappingBase): newshape: ShapeType order: str - _mapper_method: ClassVar[str] = "map_reshape" - if __debug__: - def __attrs_post_init__(self) -> None: - super().__attrs_post_init__() + def __post_init__(self) -> None: + super().__post_init__() @property def shape(self) -> ShapeType: @@ -1558,7 +1645,7 @@ def shape(self) -> ShapeType: # {{{ indexing -@attrs.frozen(eq=False, repr=False, hash=True, cache_hash=True) +@array_dataclass() class IndexBase(_SuppliedAxesAndTagsMixin, IndexRemappingBase): """ Abstract class for all index expressions on an array. @@ -1568,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: @@ -1681,7 +1768,7 @@ def shape(self) -> ShapeType: # {{{ base class for arguments -@attrs.frozen(eq=False, repr=False, hash=True, cache_hash=True) +@array_dataclass() class InputArgumentBase(Array): r"""Base class for input arguments. @@ -1721,7 +1808,7 @@ def dtype(self) -> np.dtype[Any]: pass -@attrs.frozen(eq=False, repr=False, hash=False) +@array_dataclass(hash=False) class DataWrapper(_SuppliedAxesAndTagsMixin, InputArgumentBase): """Takes concrete array data and packages it to be compatible with the :class:`Array` interface. @@ -1762,9 +1849,7 @@ class DataWrapper(_SuppliedAxesAndTagsMixin, InputArgumentBase): (i.e. the very same instance). """ data: DataInterface - _shape: ShapeType - - _mapper_method: ClassVar[str] = "map_data_wrapper" + shape: ShapeType @property def name(self) -> None: @@ -1781,10 +1866,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 @@ -1794,7 +1875,7 @@ def dtype(self) -> np.dtype[Any]: # {{{ placeholder -@attrs.frozen(eq=False, repr=False, hash=True, cache_hash=True) +@array_dataclass() class Placeholder( _SuppliedAxesAndTagsMixin, _SuppliedShapeAndDtypeMixin, @@ -1811,14 +1892,12 @@ class Placeholder( """ name: str - _mapper_method: ClassVar[str] = "map_placeholder" - # }}} # {{{ size parameter -@attrs.frozen(eq=False, repr=False, hash=True, cache_hash=True) +@array_dataclass() class SizeParam( _SuppliedAxesAndTagsMixin, InputArgumentBase): @@ -1830,10 +1909,8 @@ class SizeParam( The name by which a value is supplied for the argument once computation begins. """ - name: str - axes: AxesT = attrs.field(kw_only=True, default=()) - - _mapper_method: ClassVar[str] = "map_size_param" + name: str = dataclasses.field(kw_only=True) + axes: AxesT = dataclasses.field(kw_only=True, default=()) @property def shape(self) -> ShapeType: @@ -2176,7 +2253,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(),) @@ -2198,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. " @@ -2305,7 +2381,7 @@ def eye(N: int, M: int | None = None, k: int = 0, # noqa: N803 # {{{ arange -@attrs.define +@dataclasses.dataclass class _ArangeInfo: start: int | None stop: int | None diff --git a/pytato/codegen.py b/pytato/codegen.py index 3a1835930..aa26e48a9 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -144,7 +144,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 8c2b4c69d..c7edf7616 100644 --- a/pytato/distributed/nodes.py +++ b/pytato/distributed/nodes.py @@ -53,10 +53,10 @@ THE SOFTWARE. """ +import dataclasses from collections.abc import Hashable -from typing import Any, ClassVar +from typing import Any -import attrs import numpy as np from pytools.tag import Tag, Taggable @@ -71,6 +71,7 @@ _get_default_tags, _SuppliedAxesAndTagsMixin, _SuppliedShapeAndDtypeMixin, + array_dataclass, normalize_shape, ) @@ -80,7 +81,7 @@ # {{{ send -@attrs.frozen(init=True, eq=True, hash=True, cache_hash=True) +@array_dataclass() class DistributedSend(Taggable): """Class representing a distributed send operation. See :class:`DistributedSendRefHolder` for a way to ensure that nodes @@ -104,20 +105,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, hash=True, cache_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 @@ -155,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 @@ -182,7 +181,7 @@ def non_equality_tags(self) -> frozenset[Tag]: # {{{ receive -@attrs.frozen(eq=False, hash=True, cache_hash=True) +@array_dataclass() class DistributedRecv(_SuppliedAxesAndTagsMixin, _SuppliedShapeAndDtypeMixin, Array): """Class representing a distributed receive operation. @@ -212,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/distributed/partition.py b/pytato/distributed/partition.py index c9822549d..ca817ad85 100644 --- a/pytato/distributed/partition.py +++ b/pytato/distributed/partition.py @@ -63,6 +63,7 @@ """ import collections +import dataclasses from collections.abc import Hashable, Iterable, Iterator, Mapping, Sequence, Set from functools import reduce from typing import ( @@ -73,7 +74,6 @@ cast, ) -import attrs from immutabledict import immutabledict from pymbolic.mapper.optimize import optimize_mapper @@ -97,7 +97,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). @@ -186,7 +186,7 @@ def __sub__(self, other: Set[_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. @@ -246,7 +246,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 @@ -364,7 +364,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 a613ab838..2d7301eb8 100644 --- a/pytato/distributed/tags.py +++ b/pytato/distributed/tags.py @@ -102,7 +102,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 917978097..84d8a21ed 100644 --- a/pytato/distributed/verify.py +++ b/pytato/distributed/verify.py @@ -36,11 +36,11 @@ """ +import dataclasses import logging from collections.abc import Sequence from typing import TYPE_CHECKING, Any -import attrs import numpy as np from pymbolic.mapper.optimize import optimize_mapper @@ -69,7 +69,7 @@ # {{{ data structures -@attrs.define(frozen=True) +@dataclasses.dataclass(frozen=True) class _SummarizedDistributedSend: src_rank: int dest_rank: int @@ -79,19 +79,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 fc0a3c192..3450e62e6 100644 --- a/pytato/function.py +++ b/pytato/function.py @@ -55,17 +55,16 @@ THE SOFTWARE. """ +import dataclasses import enum import re from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping from functools import cached_property from typing import ( Any, - ClassVar, TypeVar, ) -import attrs from immutabledict import immutabledict from pytools import memoize_method @@ -78,6 +77,7 @@ Placeholder, ShapeType, _dtype_any, + array_dataclass, ) @@ -97,8 +97,7 @@ class ReturnType(enum.Enum): TUPLE_OF_ARRAYS = 2 -# eq=False to avoid equality comparison without EqualityMaper -@attrs.define(frozen=True, eq=False, hash=True, cache_hash=True) +@array_dataclass() class FunctionDefinition(Taggable): r""" A function definition that represents its outputs as instances of @@ -151,9 +150,12 @@ 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) + + if __debug__: + def __post_init__(self) -> None: + assert isinstance(self.returns, immutabledict) @cached_property def _placeholders(self) -> Mapping[str, Placeholder]: @@ -187,7 +189,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 ) -> Array | tuple[Array, ...] | Mapping[str, Array]: @@ -241,7 +243,7 @@ def __eq__(self, other: Any) -> bool: return EqualityComparer().map_function_definition(self, other) -@attrs.frozen(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`. @@ -255,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: @@ -282,16 +283,15 @@ def call(self) -> Call: @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 -@attrs.define(frozen=True, eq=False, hash=True, cache_hash=True, repr=False) +@array_dataclass() class Call(AbstractResultWithNamedArrays): """ Records an invocation to a :class:`FunctionDefinition`. @@ -307,19 +307,17 @@ class Call(AbstractResultWithNamedArrays): """ function: FunctionDefinition - bindings: Mapping[str, Array] = attrs.field( - validator=attrs.validators.instance_of(immutabledict)) - - _mapper_method: ClassVar[str] = "map_call" + bindings: Mapping[str, Array] - 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__() + assert isinstance(self.bindings, immutabledict) + super().__post_init__() def __contains__(self, name: object) -> bool: return name in self.function.returns @@ -339,7 +337,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 c316d1cf6..118979eea 100644 --- a/pytato/loopy.py +++ b/pytato/loopy.py @@ -28,14 +28,13 @@ """ +import dataclasses from collections.abc import Iterable, Iterator, Mapping, Sequence from numbers import Number from typing import ( Any, - ClassVar, ) -import attrs import islpy as isl import numpy as np from immutabledict import immutabledict @@ -52,6 +51,7 @@ NamedArray, ShapeType, SizeParam, + array_dataclass, ) from pytato.scalar_expr import ( EvaluationMapper, @@ -91,20 +91,21 @@ """ -@attrs.frozen(eq=False) +@array_dataclass() 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 = dataclasses.replace - copy = attrs.evolve + def __post_init__(self) -> None: + assert isinstance(self.bindings, immutabledict) + super().__post_init__() @property def _result_names(self) -> frozenset[str]: @@ -136,7 +137,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 @@ -151,13 +152,14 @@ def __iter__(self) -> Iterator[str]: return iter(self._result_names) -@attrs.frozen(eq=False, hash=True, cache_hash=True) -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 diff --git a/pytato/stringifier.py b/pytato/stringifier.py index a0ce7754f..6b50d67f1 100644 --- a/pytato/stringifier.py +++ b/pytato/stringifier.py @@ -25,9 +25,9 @@ THE SOFTWARE. """ +import dataclasses from typing import Any, cast -import attrs import numpy as np from immutabledict import immutabledict @@ -104,7 +104,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") @@ -165,7 +165,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 @@ -182,7 +182,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 08709bb42..c1f72f214 100644 --- a/pytato/target/loopy/codegen.py +++ b/pytato/target/loopy/codegen.py @@ -23,12 +23,12 @@ THE SOFTWARE. """ +import dataclasses import re import sys from abc import ABC, abstractmethod from collections.abc import Mapping -import attrs import islpy as isl import loopy as lp @@ -144,7 +144,7 @@ def loopy_substitute( # {{{ LoopyExpressionContexts -@attrs.define(init=True, repr=False, eq=False) +@dataclasses.dataclass(init=True, repr=False, eq=False) class PersistentExpressionContext: """ Mutable state used while generating :mod:`loopy` expressions for a @@ -168,8 +168,7 @@ class PersistentExpressionContext: """ 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]: @@ -179,7 +178,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 @@ -303,7 +302,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 @@ -327,7 +326,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`. @@ -349,10 +348,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()) @@ -1012,7 +1011,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 25edfc7da..4c5d4a0f1 100644 --- a/pytato/transform/__init__.py +++ b/pytato/transform/__init__.py @@ -26,9 +26,9 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. """ +import dataclasses import logging from collections.abc import Callable, Hashable, Iterable, Mapping -from dataclasses import dataclass from typing import ( Any, Generic, @@ -37,7 +37,6 @@ cast, ) -import attrs import numpy as np from immutabledict import immutabledict from typing_extensions import Self @@ -452,7 +451,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, @@ -491,7 +490,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), @@ -618,7 +617,7 @@ def map_data_wrapper(self, expr: DataWrapper, def map_size_param(self, expr: SizeParam, *args: P.args, **kwargs: P.kwargs) -> 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: P.args, **kwargs: P.kwargs) -> Array: return Einsum(expr.access_descriptors, @@ -665,7 +664,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, @@ -1329,7 +1328,7 @@ def rec(self, expr: ArrayOrNames) -> ArrayOrNames: # {{{ MPMS materializer -@dataclass(frozen=True, eq=True) +@dataclasses.dataclass(frozen=True, eq=True) class MPMSMaterializerAccumulator: """This class serves as the return value of :class:`MPMSMaterializer`. It contains the set of materialized predecessors and the rewritten expression diff --git a/pytato/transform/einsum_distributive_law.py b/pytato/transform/einsum_distributive_law.py index 33ae4b100..0d9a6076f 100644 --- a/pytato/transform/einsum_distributive_law.py +++ b/pytato/transform/einsum_distributive_law.py @@ -34,10 +34,10 @@ """ +import dataclasses from collections.abc import Callable, Mapping from typing import cast -import attrs import numpy as np from immutabledict import immutabledict @@ -74,7 +74,7 @@ class EinsumDistributiveLawDescriptor: """ -@attrs.frozen +@dataclasses.dataclass(frozen=True) class DoNotDistribute(EinsumDistributiveLawDescriptor): """ Tells :func:`apply_distributive_property_to_einsums` to not apply @@ -82,7 +82,7 @@ class DoNotDistribute(EinsumDistributiveLawDescriptor): """ -@attrs.frozen +@dataclasses.dataclass(frozen=True) class DoDistribute(EinsumDistributiveLawDescriptor): """ Tells :func:`apply_distributive_property_to_einsums` to apply distributive @@ -91,16 +91,16 @@ 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] - 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 75b617ebf..f505c900a 100644 --- a/pytato/visualization/dot.py +++ b/pytato/visualization/dot.py @@ -27,6 +27,7 @@ """ +import dataclasses import html from collections.abc import Callable, Mapping from functools import partial @@ -35,8 +36,6 @@ Any, ) -import attrs - from pytools import UniqueNameGenerator from pytools.codegen import remove_common_indentation from pytools.tag import Tag @@ -81,7 +80,7 @@ # {{{ _DotEmitter -@attrs.define +@dataclasses.dataclass class _SubgraphTree: contents: list[str] | None subgraphs: dict[str, _SubgraphTree] @@ -156,7 +155,7 @@ def emit_subgraph(sg: _SubgraphTree) -> None: # {{{ array -> dot node converter -@attrs.define +@dataclasses.dataclass class _DotNodeInfo: title: str fields: dict[str, Any] @@ -202,7 +201,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/test/test_pytato.py b/test/test_pytato.py index 271c8fb01..46c794d3e 100644 --- a/test/test_pytato.py +++ b/test/test_pytato.py @@ -27,9 +27,9 @@ THE SOFTWARE. """ +import dataclasses import sys -import attrs import numpy as np import pytest from testlib import RandomDAGContext, make_random_dag @@ -1189,14 +1189,14 @@ def test_with_tagged_reduction(): def test_derived_class_uses_correct_array_eq(): - @attrs.define(frozen=True) + @dataclasses.dataclass(frozen=True) class MyNewArrayT(_SuppliedAxesAndTagsMixin, 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(_SuppliedAxesAndTagsMixin, pt.Array): pass