nshutils 0.2.0__tar.gz → 0.3.0__tar.gz

{nshutils-0.2.0 → nshutils-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshutils
-Version: 0.2.0
+Version: 0.3.0
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com
@@ -11,6 +11,7 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Requires-Dist: beartype (>=0.18.5,<0.19.0)
 Requires-Dist: jaxtyping (>=0.2.33,<0.3.0)
+Requires-Dist: numpy
 Requires-Dist: pysnooper (>=1.2.0,<2.0.0)
 Requires-Dist: typing-extensions
 Description-Content-Type: text/markdown

{nshutils-0.2.0 → nshutils-0.3.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nshutils"
-version = "0.2.0"
+version = "0.3.0"
 description = ""
 authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
 readme = "README.md"
@@ -11,8 +11,21 @@ pysnooper = "^1.2.0"
 jaxtyping = "^0.2.33"
 typing-extensions = "*"
 beartype = "^0.18.5"
+numpy = "*"
 
 
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
+
+[tool.pyright]
+typeCheckingMode = "standard"
+deprecateTypingAliases = true
+strictListInference = true
+strictDictionaryInference = true
+strictSetInference = true
+reportPrivateImportUsage = false
+ignore = ["./build/"]
+
+[tool.ruff.lint]
+ignore = ["F722", "F821", "E731", "E741"]

{nshutils-0.2.0 → nshutils-0.3.0}/src/nshutils/__init__.py

@@ -1,2 +1,3 @@
+from . import actsave as actsave
 from . import typecheck as typecheck
 from .snoop import snoop as snoop

nshutils-0.3.0/src/nshutils/actsave/__init__.py

@@ -0,0 +1,6 @@
+from ._loader import ActivationLoader as ActivationLoader
+from ._loader import ActLoad as ActLoad
+from ._saver import Activation as Activation
+from ._saver import ActivationSaver as ActivationSaver
+from ._saver import ActSave as ActSave
+from ._saver import Transform as Transform

nshutils-0.3.0/src/nshutils/actsave/_loader.py

@@ -0,0 +1,144 @@
+import pprint
+from dataclasses import dataclass, field
+from functools import cached_property
+from logging import getLogger
+from pathlib import Path
+from typing import cast, overload
+
+import numpy as np
+from typing_extensions import TypeVar, override
+
+log = getLogger(__name__)
+
+T = TypeVar("T", infer_variance=True)
+
+
+@dataclass
+class LoadedActivation:
+    base_dir: Path = field(repr=False)
+    name: str
+    num_activations: int = field(init=False)
+    activation_files: list[Path] = field(init=False, repr=False)
+
+    def __post_init__(self):
+        if not self.activation_dir.exists():
+            raise ValueError(f"Activation dir {self.activation_dir} does not exist")
+
+        # The number of activations = the * of .npy files in the activation dir
+        self.activation_files = list(self.activation_dir.glob("*.npy"))
+        # Sort the activation files by the numerical index in the filename
+        self.activation_files.sort(key=lambda p: int(p.stem))
+        self.num_activations = len(self.activation_files)
+
+    @property
+    def activation_dir(self) -> Path:
+        return self.base_dir / self.name
+
+    def _load_activation(self, item: int):
+        activation_path = self.activation_files[item]
+        if not activation_path.exists():
+            raise ValueError(f"Activation {activation_path} does not exist")
+        return cast(np.ndarray, np.load(activation_path, allow_pickle=True))
+
+    @overload
+    def __getitem__(self, item: int) -> np.ndarray: ...
+
+    @overload
+    def __getitem__(self, item: slice | list[int]) -> list[np.ndarray]: ...
+
+    def __getitem__(
+        self, item: int | slice | list[int]
+    ) -> np.ndarray | list[np.ndarray]:
+        if isinstance(item, int):
+            return self._load_activation(item)
+        elif isinstance(item, slice):
+            return [
+                self._load_activation(i)
+                for i in range(*item.indices(self.num_activations))
+            ]
+        elif isinstance(item, list):
+            return [self._load_activation(i) for i in item]
+        else:
+            raise TypeError(f"Invalid type {type(item)} for item {item}")
+
+    def __iter__(self):
+        return iter(self[i] for i in range(self.num_activations))
+
+    def __len__(self):
+        return self.num_activations
+
+    def all_activations(self):
+        return [self[i] for i in range(self.num_activations)]
+
+    @override
+    def __repr__(self):
+        return f"<LoadedActivation {self.name} ({self.num_activations} activations)>"
+
+
+class ActLoad:
+    @classmethod
+    def all_versions(cls, dir: str | Path):
+        dir = Path(dir)
+
+        # If the dir is not an activation base directory, we return None
+        if not (dir / ".activationbase").exists():
+            return None
+
+        # The contents of `dir` should be directories, each of which is a version.
+        return [
+            (subdir, int(subdir.name)) for subdir in dir.iterdir() if subdir.is_dir()
+        ]
+
+    @classmethod
+    def is_valid_activation_base(cls, dir: str | Path):
+        return cls.all_versions(dir) is not None
+
+    @classmethod
+    def from_latest_version(cls, dir: str | Path):
+        # The contents of `dir` should be directories, each of which is a version
+        # We need to find the latest version
+        if (all_versions := cls.all_versions(dir)) is None:
+            raise ValueError(f"{dir} is not an activation base directory")
+
+        path, _ = max(all_versions, key=lambda p: p[1])
+        return cls(path)
+
+    def __init__(self, dir: Path):
+        self._dir = dir
+
+    def activation(self, name: str):
+        return LoadedActivation(self._dir, name)
+
+    @cached_property
+    def activations(self):
+        dirs = list(self._dir.iterdir())
+        # Sort the dirs by the last modified time
+        dirs.sort(key=lambda p: p.stat().st_mtime)
+
+        return {p.name: LoadedActivation(self._dir, p.name) for p in dirs}
+
+    def __iter__(self):
+        return iter(self.activations.values())
+
+    def __getitem__(self, item: str):
+        return self.activations[item]
+
+    def __len__(self):
+        return len(self.activations)
+
+    @override
+    def __repr__(self):
+        acts_str = pprint.pformat(
+            {
+                name: f"<{activation.num_activations} activations>"
+                for name, activation in self.activations.items()
+            }
+        )
+        acts_str = acts_str.replace("'<", "<").replace(">'", ">")
+        return f"ActLoad({acts_str})"
+
+    def get(self, name: str, /, default: T) -> LoadedActivation | T:
+        return self.activations.get(name, default)
+
+
+ActivationLoader = ActLoad

nshutils-0.3.0/src/nshutils/actsave/_saver.py

@@ -0,0 +1,360 @@
+import contextlib
+import fnmatch
+import tempfile
+import uuid
+import weakref
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from functools import wraps
+from logging import getLogger
+from pathlib import Path
+from typing import TYPE_CHECKING, Generic, TypeAlias, cast, overload
+
+import numpy as np
+from typing_extensions import Never, ParamSpec, TypeVar, override
+
+from ..collections import apply_to_collection
+
+try:
+    import torch
+
+    if not TYPE_CHECKING:
+        Tensor: TypeAlias = torch.Tensor
+except ImportError:
+    torch = None
+
+    if not TYPE_CHECKING:
+        Tensor: TypeAlias = Never
+
+if TYPE_CHECKING:
+    Tensor: TypeAlias = Never
+
+log = getLogger(__name__)
+
+Value: TypeAlias = int | float | complex | bool | str | np.ndarray | Tensor | None
+ValueOrLambda = Value | Callable[..., Value]
+
+
+def _torch_is_scripting() -> bool:
+    if torch is None:
+        return False
+
+    return torch.jit.is_scripting()
+
+
+def _to_numpy(activation: Value) -> np.ndarray:
+    # Make sure it's not `None`
+    if activation is None:
+        raise ValueError("Activation should not be `None`")
+
+    if isinstance(activation, (int, float, complex, str, bool)):
+        return np.array(activation)
+    elif isinstance(activation, np.ndarray):
+        return activation
+    elif isinstance(activation, Tensor):
+        activation = activation.detach()
+        if activation.is_floating_point():
+            # NOTE: We need to convert to float32 because [b]float16 is not supported by numpy
+            activation = activation.float()
+        return activation.cpu().numpy()
+    else:
+        log.warning(f"Unrecognized activation type {type(activation)}")
+
+    return activation
+
+
+T = TypeVar("T", infer_variance=True)
+
+
+# A wrapper around weakref.ref that allows for primitive types
+# To get around errors like:
+# TypeError: cannot create weak reference to 'int' object
+class WeakRef(Generic[T]):
+    _ref: Callable[[], T] | None
+
+    def __init__(self, obj: T):
+        try:
+            self._ref = cast(Callable[[], T], weakref.ref(obj))
+        except TypeError as e:
+            if "cannot create weak reference" not in str(e):
+                raise
+            self._ref = lambda: obj
+
+    def __call__(self) -> T:
+        if self._ref is None:
+            raise RuntimeError("WeakRef is deleted")
+        return self._ref()
+
+    def delete(self):
+        del self._ref
+        self._ref = None
+
+
+@dataclass
+class Activation:
+    name: str
+    ref: WeakRef[ValueOrLambda] | None
+    transformed: np.ndarray | None = None
+
+    def __post_init__(self):
+        # Update the `name` to replace `/` with `.`
+        self.name = self.name.replace("/", ".")
+
+    def __call__(self) -> np.ndarray | None:
+        # If we have a transformed value, we return it
+        if self.transformed is not None:
+            return self.transformed
+
+        if self.ref is None:
+            raise RuntimeError("Activation is deleted")
+
+        # If we have a lambda, we need to call it
+        unrwapped_ref = self.ref()
+        activation = unrwapped_ref
+        if callable(unrwapped_ref):
+            activation = unrwapped_ref()
+
+        # If we have a `None`, we return early
+        if activation is None:
+            return None
+
+        activation = apply_to_collection(activation, Tensor, _to_numpy)
+        activation = _to_numpy(activation)
+
+        # Set the transformed value
+        self.transformed = activation
+
+        # Delete the reference
+        self.ref.delete()
+        del self.ref
+        self.ref = None
+
+        return self.transformed
+
+    @classmethod
+    def from_value_or_lambda(cls, name: str, value_or_lambda: ValueOrLambda):
+        return cls(name, WeakRef(value_or_lambda))
+
+    @classmethod
+    def from_dict(cls, d: Mapping[str, ValueOrLambda]):
+        return [cls.from_value_or_lambda(k, v) for k, v in d.items()]
+
+
+Transform = Callable[[Activation], Mapping[str, ValueOrLambda]]
+
+
+def _ensure_supported():
+    try:
+        import torch.distributed as dist
+
+        if dist.is_initialized() and dist.get_world_size() > 1:
+            raise RuntimeError("Only single GPU is supported at the moment")
+    except ImportError:
+        pass
+
+
+P = ParamSpec("P")
+
+
+def _ignore_if_scripting(fn: Callable[P, None]) -> Callable[P, None]:
+    @wraps(fn)
+    def wrapper(*args: P.args, **kwargs: P.kwargs) -> None:
+        if _torch_is_scripting():
+            return
+
+        _ensure_supported()
+        fn(*args, **kwargs)
+
+    return wrapper
+
+
+class _Saver:
+    def __init__(
+        self,
+        save_dir: Path,
+        prefixes_fn: Callable[[], list[str]],
+        *,
+        filters: list[str] | None = None,
+    ):
+        # Create a directory under `save_dir` by autoincrementing
+        # (i.e., every activation save context, we create a new directory)
+        # The id = the number of activation subdirectories
+        self._id = sum(1 for subdir in save_dir.glob("*") if subdir.is_dir())
+        save_dir.mkdir(parents=True, exist_ok=True)
+
+        # Add a .activationbase file to the save_dir to indicate that this is an activation base
+        (save_dir / ".activationbase").touch(exist_ok=True)
+
+        self._save_dir = save_dir / f"{self._id:04d}"
+        # Make sure `self._save_dir` does not exist and create it
+        self._save_dir.mkdir(exist_ok=False)
+
+        self._prefixes_fn = prefixes_fn
+        self._filters = filters
+
+    def _save_activation(self, activation: Activation):
+        # If the activation value is `None`, we skip it.
+        if (activation_value := activation()) is None:
+            return
+
+        # Save the activation to self._save_dir / name / {id}.npz, where id is an auto-incrementing integer
+        file_name = ".".join(self._prefixes_fn() + [activation.name])
+        path = self._save_dir / file_name
+        path.mkdir(exist_ok=True, parents=True)
+
+        # Get the next id and save the activation
+        id = len(list(path.glob("*.npy")))
+        np.save(path / f"{id:04d}.npy", activation_value)
+
+    @_ignore_if_scripting
+    def save(
+        self,
+        acts: dict[str, ValueOrLambda] | None = None,
+        /,
+        **kwargs: ValueOrLambda,
+    ):
+        kwargs.update(acts or {})
+
+        # Build activations
+        activations = Activation.from_dict(kwargs)
+
+        for activation in activations:
+            # Make sure name matches at least one filter if filters are specified
+            if self._filters is not None and all(
+                not fnmatch.fnmatch(activation.name, f) for f in self._filters
+            ):
+                continue
+
+            # Save the current activation
+            self._save_activation(activation)
+
+        del activations
+
+
+class ActSaveProvider:
+    _saver: _Saver | None = None
+    _prefixes: list[str] = []
+
+    def initialize(self, save_dir: Path | None = None):
+        """
+        Initializes the saver with the given configuration and save directory.
+
+        Args:
+            save_dir (Path): The directory where the saved files will be stored.
+        """
+        if self._saver is None:
+            if save_dir is None:
+                save_dir = Path(tempfile.gettempdir()) / f"actsave-{uuid.uuid4()}"
+                log.critical(f"No save_dir specified, using {save_dir=}")
+            self._saver = _Saver(
+                save_dir,
+                lambda: self._prefixes,
+            )
+
+    @contextlib.contextmanager
+    def enabled(self, save_dir: Path | None = None):
+        """
+        Context manager that enables the actsave functionality with the specified configuration.
+
+        Args:
+            save_dir (Path): The directory where the saved files will be stored.
+        """
+        prev = self._saver
+        self.initialize(save_dir)
+        try:
+            yield
+        finally:
+            self._saver = prev
+
+    @override
+    def __init__(self):
+        super().__init__()
+
+        self._saver = None
+        self._prefixes = []
+
+    @contextlib.contextmanager
+    def context(self, label: str):
+        """
+        A context manager that adds a label to the current context.
+
+        Args:
+            label (str): The label for the context.
+        """
+        if _torch_is_scripting():
+            yield
+            return
+
+        if self._saver is None:
+            yield
+            return
+
+        _ensure_supported()
+
+        log.debug(f"Entering ActSave context {label}")
+        self._prefixes.append(label)
+        try:
+            yield
+        finally:
+            _ = self._prefixes.pop()
+
+    prefix = context
+
+    @overload
+    def __call__(
+        self,
+        acts: dict[str, ValueOrLambda] | None = None,
+        /,
+        **kwargs: ValueOrLambda,
+    ):
+        """
+        Saves the activations to disk.
+
+        Args:
+            acts (dict[str, ValueOrLambda] | None, optional): A dictionary of acts. Defaults to None.
+            **kwargs (ValueOrLambda): Additional keyword arguments.
+
+        Returns:
+            None
+
+        """
+        ...
+
+    @overload
+    def __call__(self, acts: Callable[[], dict[str, ValueOrLambda]], /):
+        """
+        Saves the activations to disk.
+
+        Args:
+            acts (Callable[[], dict[str, ValueOrLambda]]): A callable that returns a dictionary of acts.
+            **kwargs (ValueOrLambda): Additional keyword arguments.
+
+        Returns:
+            None
+
+        """
+        ...
+
+    def __call__(
+        self,
+        acts: (
+            dict[str, ValueOrLambda] | Callable[[], dict[str, ValueOrLambda]] | None
+        ) = None,
+        /,
+        **kwargs: ValueOrLambda,
+    ):
+        if _torch_is_scripting():
+            return
+
+        if self._saver is None:
+            return
+
+        if acts is not None and callable(acts):
+            acts = acts()
+        self._saver.save(acts, **kwargs)
+
+    save = __call__
+
+
+ActSave = ActSaveProvider()
+ActivationSaver = ActSave

nshutils-0.3.0/src/nshutils/collections.py

@@ -0,0 +1,271 @@
+# Copyright The PyTorch Lightning team.
+# Licensed under the Apache License, Version 2.0 (the "License");
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+import dataclasses
+from collections import OrderedDict, defaultdict
+from collections.abc import Callable, Mapping, Sequence
+from copy import deepcopy
+from typing import Any
+
+
+def is_namedtuple(obj: object) -> bool:
+    """Check if object is type nametuple."""
+    # https://github.com/pytorch/pytorch/blob/v1.8.1/torch/nn/parallel/scatter_gather.py#L4-L8
+    return (
+        isinstance(obj, tuple) and hasattr(obj, "_asdict") and hasattr(obj, "_fields")
+    )
+
+
+def is_dataclass_instance(obj: object) -> bool:
+    """Check if object is dataclass."""
+    # https://docs.python.org/3/library/dataclasses.html#module-level-decorators-classes-and-functions
+    return dataclasses.is_dataclass(obj) and not isinstance(obj, type)
+
+
+def apply_to_collection(
+    data: Any,
+    dtype: type | Any | tuple[type | Any],
+    function: Callable,
+    *args: Any,
+    wrong_dtype: type | tuple[type, ...] | None = None,
+    include_none: bool = True,
+    allow_frozen: bool = False,
+    **kwargs: Any,
+) -> Any:
+    """Recursively applies a function to all elements of a certain dtype.
+
+    Args:
+        data: the collection to apply the function to
+        dtype: the given function will be applied to all elements of this dtype
+        function: the function to apply
+        *args: positional arguments (will be forwarded to calls of ``function``)
+        wrong_dtype: the given function won't be applied if this type is specified and the given collections
+            is of the ``wrong_dtype`` even if it is of type ``dtype``
+        include_none: Whether to include an element if the output of ``function`` is ``None``.
+        allow_frozen: Whether not to error upon encountering a frozen dataclass instance.
+        **kwargs: keyword arguments (will be forwarded to calls of ``function``)
+
+    Returns:
+        The resulting collection
+    """
+    # Breaking condition
+    if isinstance(data, dtype) and (
+        wrong_dtype is None or not isinstance(data, wrong_dtype)
+    ):
+        return function(data, *args, **kwargs)
+
+    elem_type = type(data)
+
+    # Recursively apply to collection items
+    if isinstance(data, Mapping):
+        out = []
+        for k, v in data.items():
+            v = apply_to_collection(
+                v,
+                dtype,
+                function,
+                *args,
+                wrong_dtype=wrong_dtype,
+                include_none=include_none,
+                allow_frozen=allow_frozen,
+                **kwargs,
+            )
+            if include_none or v is not None:
+                out.append((k, v))
+        if isinstance(data, defaultdict):
+            return elem_type(data.default_factory, OrderedDict(out))
+        return elem_type(OrderedDict(out))
+
+    is_namedtuple_ = is_namedtuple(data)
+    is_sequence = isinstance(data, Sequence) and not isinstance(data, str)
+    if is_namedtuple_ or is_sequence:
+        out = []
+        for d in data:
+            v = apply_to_collection(
+                d,
+                dtype,
+                function,
+                *args,
+                wrong_dtype=wrong_dtype,
+                include_none=include_none,
+                allow_frozen=allow_frozen,
+                **kwargs,
+            )
+            if include_none or v is not None:
+                out.append(v)
+        return elem_type(*out) if is_namedtuple_ else elem_type(out)
+
+    if is_dataclass_instance(data):
+        # make a deepcopy of the data,
+        # but do not deepcopy mapped fields since the computation would
+        # be wasted on values that likely get immediately overwritten
+        fields = {}
+        memo = {}
+        for field in dataclasses.fields(data):
+            field_value = getattr(data, field.name)
+            fields[field.name] = (field_value, field.init)
+            memo[id(field_value)] = field_value
+        result = deepcopy(data, memo=memo)
+        # apply function to each field
+        for field_name, (field_value, field_init) in fields.items():
+            v = None
+            if field_init:
+                v = apply_to_collection(
+                    field_value,
+                    dtype,
+                    function,
+                    *args,
+                    wrong_dtype=wrong_dtype,
+                    include_none=include_none,
+                    allow_frozen=allow_frozen,
+                    **kwargs,
+                )
+            if not field_init or (not include_none and v is None):  # retain old value
+                v = getattr(data, field_name)
+            try:
+                setattr(result, field_name, v)
+            except dataclasses.FrozenInstanceError as e:
+                if allow_frozen:
+                    # Quit early if we encounter a frozen data class; return `result` as is.
+                    break
+                raise ValueError(
+                    "A frozen dataclass was passed to `apply_to_collection` but this is not allowed."
+                ) from e
+        return result
+
+    # data is neither of dtype, nor a collection
+    return data
+
+
+def apply_to_collections(
+    data1: Any | None,
+    data2: Any | None,
+    dtype: type | Any | tuple[type | Any],
+    function: Callable,
+    *args: Any,
+    wrong_dtype: type | tuple[type] | None = None,
+    **kwargs: Any,
+) -> Any:
+    """Zips two collections and applies a function to their items of a certain dtype.
+
+    Args:
+        data1: The first collection
+        data2: The second collection
+        dtype: the given function will be applied to all elements of this dtype
+        function: the function to apply
+        *args: positional arguments (will be forwarded to calls of ``function``)
+        wrong_dtype: the given function won't be applied if this type is specified and the given collections
+            is of the ``wrong_dtype`` even if it is of type ``dtype``
+        **kwargs: keyword arguments (will be forwarded to calls of ``function``)
+
+    Returns:
+        The resulting collection
+
+    Raises:
+        AssertionError:
+            If sequence collections have different data sizes.
+    """
+    if data1 is None:
+        if data2 is None:
+            return None
+        # in case they were passed reversed
+        data1, data2 = data2, None
+
+    elem_type = type(data1)
+
+    if (
+        isinstance(data1, dtype)
+        and data2 is not None
+        and (wrong_dtype is None or not isinstance(data1, wrong_dtype))
+    ):
+        return function(data1, data2, *args, **kwargs)
+
+    if isinstance(data1, Mapping) and data2 is not None:
+        # use union because we want to fail if a key does not exist in both
+        zipped = {k: (data1[k], data2[k]) for k in data1.keys() | data2.keys()}
+        return elem_type(
+            {
+                k: apply_to_collections(
+                    *v, dtype, function, *args, wrong_dtype=wrong_dtype, **kwargs
+                )
+                for k, v in zipped.items()
+            }
+        )
+
+    is_namedtuple_ = is_namedtuple(data1)
+    is_sequence = isinstance(data1, Sequence) and not isinstance(data1, str)
+    if (is_namedtuple_ or is_sequence) and data2 is not None:
+        if len(data1) != len(data2):
+            raise ValueError("Sequence collections have different sizes.")
+        out = [
+            apply_to_collections(
+                v1, v2, dtype, function, *args, wrong_dtype=wrong_dtype, **kwargs
+            )
+            for v1, v2 in zip(data1, data2)
+        ]
+        return elem_type(*out) if is_namedtuple_ else elem_type(out)
+
+    if is_dataclass_instance(data1) and data2 is not None:
+        if not is_dataclass_instance(data2):
+            raise TypeError(
+                "Expected inputs to be dataclasses of the same type or to have identical fields"
+                f" but got input 1 of type {type(data1)} and input 2 of type {type(data2)}."
+            )
+        if not (
+            len(dataclasses.fields(data1)) == len(dataclasses.fields(data2))
+            and all(
+                map(
+                    lambda f1, f2: isinstance(f1, type(f2)),
+                    dataclasses.fields(data1),
+                    dataclasses.fields(data2),
+                )
+            )
+        ):
+            raise TypeError("Dataclasses fields do not match.")
+        # make a deepcopy of the data,
+        # but do not deepcopy mapped fields since the computation would
+        # be wasted on values that likely get immediately overwritten
+        data = [data1, data2]
+        fields: list[dict] = [{}, {}]
+        memo: dict = {}
+        for i in range(len(data)):
+            for field in dataclasses.fields(data[i]):
+                field_value = getattr(data[i], field.name)
+                fields[i][field.name] = (field_value, field.init)
+                if i == 0:
+                    memo[id(field_value)] = field_value
+
+        result = deepcopy(data1, memo=memo)
+
+        # apply function to each field
+        for (field_name, (field_value1, field_init1)), (
+            _,
+            (field_value2, field_init2),
+        ) in zip(fields[0].items(), fields[1].items()):
+            v = None
+            if field_init1 and field_init2:
+                v = apply_to_collections(
+                    field_value1,
+                    field_value2,
+                    dtype,
+                    function,
+                    *args,
+                    wrong_dtype=wrong_dtype,
+                    **kwargs,
+                )
+            if not field_init1 or not field_init2 or v is None:  # retain old value
+                return apply_to_collection(
+                    data1, dtype, function, *args, wrong_dtype=wrong_dtype, **kwargs
+                )
+            try:
+                setattr(result, field_name, v)
+            except dataclasses.FrozenInstanceError as e:
+                raise ValueError(
+                    "A frozen dataclass was passed to `apply_to_collections` but this is not allowed."
+                ) from e
+        return result
+
+    return apply_to_collection(
+        data1, dtype, function, *args, wrong_dtype=wrong_dtype, **kwargs
+    )