Support registered ensemble tensor buffers in processors

sdm/nn/_buffer.py

@@ -3,12 +3,14 @@
 
 from __future__ import annotations
 
-from collections.abc import Iterable, Iterator
-from typing import Any, Generic, TypeVar, cast
+from collections.abc import Callable, Iterable, Iterator, Sequence
+from typing import Any, Generic, Self, TypeVar, cast
 
 import torch
 from torch import Tensor
 
+from sdm import EnsembleTensor
+
 T = TypeVar("T", bound="Tensor | BufferList[Any]")
 
 
@@ -89,3 +91,91 @@ def set_extra_state(self, state: object) -> None:
             name = str(index)
             if is_nested and name not in self._modules:
                 self.add_module(name, BufferList())
+
+
+class _EnsembleTensorBuffer(torch.nn.Module, EnsembleTensor[Tensor]):
+    """Ensemble tensor registered as PyTorch module state."""
+
+    def __init__(
+        self,
+        groups: Sequence[Tensor],
+        locations: Sequence[tuple[int, int]],
+        *,
+        persistent: bool = True,
+    ) -> None:
+        super().__init__()
+        self._groups = tuple(groups)
+        self._locations = tuple(locations)
+        self._persistent = persistent
+        for index, group in enumerate(groups):
+            self.register_buffer(str(index), group, persistent=persistent)
+
+    def _apply(
+        self,
+        fn: Callable[[Tensor], Tensor],
+        recurse: bool = True,
+    ) -> Self:
+        super()._apply(fn, recurse=recurse)
+        self._groups = tuple(
+            cast(Tensor, group) for group in self._buffers.values()
+        )
+        return self
+
+    def get_extra_state(self) -> tuple[int, tuple[tuple[int, int], ...]]:
+        return len(self._groups), self._locations
+
+    def set_extra_state(self, state: object) -> None:
+        _, self._locations = cast(
+            tuple[int, tuple[tuple[int, int], ...]], state
+        )
+
+    def _save_to_state_dict(
+        self,
+        destination: dict[str, Any],
+        prefix: str,
+        keep_vars: bool,
+    ) -> None:
+        if self._persistent:
+            super()._save_to_state_dict(destination, prefix, keep_vars)
+
+    def _load_from_state_dict(
+        self,
+        state_dict: dict[str, Any],
+        prefix: str,
+        local_metadata: dict[str, Any],
+        strict: bool,
+        missing_keys: list[str],
+        unexpected_keys: list[str],
+        error_msgs: list[str],
+    ) -> None:
+        key = f"{prefix}_extra_state"
+        if self._persistent and key in state_dict:
+            num_groups, _ = state_dict[key]
+            previous = self._buffers.copy()
+            self._buffers.clear()
+            for index in range(num_groups):
+                name = str(index)
+                saved = state_dict.get(f"{prefix}{name}")
+                buffer = previous.get(name)
+                if isinstance(saved, Tensor):
+                    if buffer is None:
+                        buffer = torch.empty_like(saved)
+                    elif buffer.shape != saved.shape:
+                        buffer = buffer.new_empty(saved.shape)
+                elif buffer is None:
+                    buffer = torch.empty(0)
+                self.register_buffer(name, buffer)
+        super()._load_from_state_dict(
+            state_dict,
+            prefix,
+            local_metadata,
+            strict,
+            missing_keys,
+            unexpected_keys,
+            error_msgs,
+        )
+        if not self._persistent and key in missing_keys:
+            missing_keys.remove(key)
+        self._groups = tuple(
+            cast(Tensor, group) for group in self._buffers.values()
+        )

sdm/processing/common/shuffle.py

@@ -6,8 +6,7 @@
 import torch
 from torch import Tensor
 
-from sdm import EnsembleTable, Stype, TableTensor
-from sdm.nn._buffer import BufferList
+from sdm import EnsembleTable, EnsembleTensor, Stype, TableTensor
 from sdm.processing import EnsembleInvertibleMixin, EnsembleProcessor
 
 
@@ -26,13 +25,15 @@ class ShuffleColumns(EnsembleProcessor, EnsembleInvertibleMixin):
     handles_stypes = (Stype.numerical,)
     requires_fit = True
 
+    _permutations: EnsembleTensor[Tensor]
+
     def __init__(
         self,
         method: Literal["random", "latin"] = "random",
     ) -> None:
         super().__init__()
         self.method = method
-        self._permutations: BufferList[Tensor] = BufferList()
+        self.register_buffer("_permutations", None)
         self._host_permutations: tuple[tuple[int, ...], ...] = ()
 
     def get_extra_state(self) -> tuple[tuple[int, ...], ...]:
@@ -92,7 +93,13 @@ def _fit_ensemble(
                     )
                 permutations.append(permutation)
 
-        self._permutations = BufferList(permutations)
+        self._set_buffer(
+            "_permutations",
+            EnsembleTensor(
+                groups=tuple(row.unsqueeze(0) for row in permutations),
+                locations=tuple((i, 0) for i in range(len(permutations))),
+            ),
+        )
         # TODO: Add to EnsembleTable directly.
         member_ids_by_group: list[list[int]] = [
             [] for _ in range(ensemble_table.num_groups)

sdm/processing/ensemble.py

@@ -5,11 +5,13 @@
 
 import abc
 from collections.abc import Iterable
-from typing import Self
+from typing import Any, Self, cast
 
 import torch
+from torch import Tensor
 
-from sdm import EnsembleTable, Stype, TableTensor
+from sdm import EnsembleTable, EnsembleTensor, Stype, TableTensor
+from sdm.nn._buffer import _EnsembleTensorBuffer
 from sdm.processing import InvertibleMixin, Processor
 
 
@@ -31,6 +33,100 @@ class EnsembleProcessor(Processor):
     with one member.
     """
 
+    def register_buffer(
+        self,
+        name: str,
+        tensor: Tensor | EnsembleTensor[Any] | None,
+        persistent: bool = True,
+    ) -> None:
+        """Register tensor or ensemble tensor state.
+
+        Args:
+            name: Attribute name of the buffer.
+            tensor: Tensor or ensemble of tensors supporting PyTorch buffer
+                operations, or ``None`` for an uninitialized buffer.
+            persistent: Whether to include the buffer in ``state_dict``.
+        """
+        if isinstance(tensor, EnsembleTensor):
+            tensor = cast("EnsembleTensor[Tensor]", tensor)
+            self._buffers.pop(name, None)
+            self.add_module(
+                name,
+                _EnsembleTensorBuffer(
+                    groups=tensor._groups,
+                    locations=tensor._locations,
+                    persistent=persistent,
+                ),
+            )
+            if persistent:
+                self._non_persistent_buffers_set.discard(name)
+            else:
+                self._non_persistent_buffers_set.add(name)
+            return
+        if isinstance(self._modules.get(name), _EnsembleTensorBuffer):
+            del self._modules[name]
+        super().register_buffer(name, tensor, persistent=persistent)
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        module = self.__dict__.get("_modules", {}).get(name)
+        if isinstance(module, _EnsembleTensorBuffer) and (
+            value is None or isinstance(value, (Tensor, EnsembleTensor))
+        ):
+            self.register_buffer(name, value, persistent=module._persistent)
+            return
+        if isinstance(value, EnsembleTensor) and name in self.__dict__.get(
+            "_buffers", {}
+        ):
+            self.register_buffer(
+                name,
+                value,
+                persistent=name not in self._non_persistent_buffers_set,
+            )
+            return
+        super().__setattr__(name, value)
+
+    def _set_buffer(
+        self,
+        name: str,
+        value: Tensor | EnsembleTensor[Any],
+    ) -> None:
+        self.register_buffer(
+            name,
+            value,
+            persistent=name not in self._non_persistent_buffers_set,
+        )
+
+    def _load_from_state_dict(
+        self,
+        state_dict: dict[str, Any],
+        prefix: str,
+        local_metadata: dict[str, Any],
+        strict: bool,
+        missing_keys: list[str],
+        unexpected_keys: list[str],
+        error_msgs: list[str],
+    ) -> None:
+        # Materialize buffers whose type was unknown before fitting.
+        for name, buffer in tuple(self._buffers.items()):
+            if buffer is not None or name in self._non_persistent_buffers_set:
+                continue
+            state = state_dict.get(f"{prefix}{name}")
+            if isinstance(state, Tensor):
+                self.register_buffer(name, torch.empty_like(state))
+            elif f"{prefix}{name}._extra_state" in state_dict:
+                self.register_buffer(
+                    name, EnsembleTensor(groups=(), locations=())
+                )
+        super()._load_from_state_dict(
+            state_dict,
+            prefix,
+            local_metadata,
+            strict,
+            missing_keys,
+            unexpected_keys,
+            error_msgs,
+        )
+
     @staticmethod
     def as_processor(processor: object) -> EnsembleProcessor:
         r"""Normalize a processor-like object to a :class:`EnsembleProcessor`.

sdm/processing/numerical/detrend.py

@@ -4,8 +4,7 @@
 import torch
 from torch import Tensor
 
-from sdm import EnsembleTable, Stype
-from sdm.nn._buffer import BufferList
+from sdm import EnsembleTable, EnsembleTensor, Stype
 from sdm.processing import EnsembleInvertibleMixin, SeriesEnsembleProcessor
 from sdm.processing.numerical._stats import _isfinite
 
@@ -22,11 +21,14 @@ class Detrend(SeriesEnsembleProcessor, EnsembleInvertibleMixin):
     handles_stypes = (Stype.numerical,)
     requires_fit = True
 
+    _slopes: EnsembleTensor[Tensor]
+    _intercepts: EnsembleTensor[Tensor]
+
     def __init__(self, threshold: float | None = None) -> None:
         super().__init__()
         self.threshold = threshold
-        self._slopes: BufferList[Tensor] = BufferList()
-        self._intercepts: BufferList[Tensor] = BufferList()
+        self.register_buffer("_slopes", None)
+        self.register_buffer("_intercepts", None)
 
     def _fit_ensemble(
         self,
@@ -66,8 +68,16 @@ def _fit_ensemble(
             slopes.append(slope)
             intercepts.append(mean - slope * time_mean)
 
-        self._slopes = BufferList(slopes)
-        self._intercepts = BufferList(intercepts)
+        self._set_buffer(
+            "_slopes",
+            EnsembleTensor(groups=slopes, locations=ensemble_table._locations),
+        )
+        self._set_buffer(
+            "_intercepts",
+            EnsembleTensor(
+                groups=intercepts, locations=ensemble_table._locations
+            ),
+        )
 
     def _transform_ensemble(
         self,
@@ -90,7 +100,10 @@ def _transform_ensemble(
             x = group.numerical
             time = torch.arange(x.size(-2), device=x.device, dtype=x.dtype)
             time = time.view(-1, 1) + self._offset[first_member_id]
-            trend = self._intercepts[group_id] + self._slopes[group_id] * time
+            trend = (
+                self._intercepts._groups[group_id]
+                + self._slopes._groups[group_id] * time
+            )
 
             groups.append(group.replace_blocks(numerical=x - trend))
 
@@ -117,9 +130,12 @@ def _inverse_transform_ensemble(
             x = group.numerical
             time = torch.arange(x.size(-2), device=x.device, dtype=x.dtype)
             time = time.view(-1, 1) + self._offset[first_member_id]
-            trend = self._intercepts[group_id] + self._slopes[group_id] * time
+            trend = (
+                self._intercepts._groups[group_id]
+                + self._slopes._groups[group_id] * time
+            )
 
-            x = x.unflatten(-1, (self._slopes[group_id].size(-1), -1))
+            x = x.unflatten(-1, (self._slopes._groups[group_id].size(-1), -1))
             x = (x + trend.unsqueeze(-1)).flatten(-2)
 
             groups.append(group.replace_blocks(numerical=x))

sdm/processing/numerical/flip_sign.py

@@ -4,8 +4,7 @@
 import torch
 from torch import Tensor
 
-from sdm import EnsembleTable, Stype
-from sdm.nn._buffer import BufferList
+from sdm import EnsembleTable, EnsembleTensor, Stype
 from sdm.processing import EnsembleInvertibleMixin, EnsembleProcessor
 
 
@@ -19,13 +18,15 @@ class FlipSign(EnsembleProcessor, EnsembleInvertibleMixin):
     handles_stypes = (Stype.numerical,)
     requires_fit = True
 
+    _signs: EnsembleTensor[Tensor]
+
     def __init__(
         self,
         probability: float = 0.5,
     ) -> None:
         super().__init__()
         self.probability = probability
-        self._signs: BufferList[Tensor] = BufferList()
+        self.register_buffer("_signs", None)
 
     def _fit_ensemble(
         self,
@@ -42,7 +43,13 @@ def _fit_ensemble(
             sign.bernoulli_(self.probability, generator=generator)
             sign.mul_(-2).add_(1)
             signs.append(sign)
-        self._signs = BufferList(signs)
+        self._set_buffer(
+            "_signs",
+            EnsembleTensor(
+                groups=tuple(sign.unsqueeze(0) for sign in signs),
+                locations=tuple((i, 0) for i in range(len(signs))),
+            ),
+        )
 
     def _transform_ensemble(
         self,