diff --git a/sdm/nn/_buffer.py b/sdm/nn/_buffer.py index 38e07bc6c..c0ca70b5d 100644 --- a/sdm/nn/_buffer.py +++ b/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() + ) diff --git a/sdm/processing/common/shuffle.py b/sdm/processing/common/shuffle.py index 856641ef4..9a66172a7 100644 --- a/sdm/processing/common/shuffle.py +++ b/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) diff --git a/sdm/processing/ensemble.py b/sdm/processing/ensemble.py index 9ece89286..316403dd5 100644 --- a/sdm/processing/ensemble.py +++ b/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`. diff --git a/sdm/processing/numerical/detrend.py b/sdm/processing/numerical/detrend.py index a3d8eaaac..0f9daf097 100644 --- a/sdm/processing/numerical/detrend.py +++ b/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)) diff --git a/sdm/processing/numerical/flip_sign.py b/sdm/processing/numerical/flip_sign.py index 57cec8958..30fb823af 100644 --- a/sdm/processing/numerical/flip_sign.py +++ b/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, diff --git a/sdm/processing/numerical/interpolate.py b/sdm/processing/numerical/interpolate.py index 540f5ed1a..ded90301f 100644 --- a/sdm/processing/numerical/interpolate.py +++ b/sdm/processing/numerical/interpolate.py @@ -5,8 +5,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 @@ -22,6 +21,9 @@ class Interpolate(SeriesEnsembleProcessor, EnsembleInvertibleMixin): handles_stypes = (Stype.numerical,) requires_fit = True + _last_values: EnsembleTensor[Tensor] + _last_indices: EnsembleTensor[Tensor] + def __init__( self, *, @@ -31,8 +33,8 @@ def __init__( super().__init__() self.fill_value = fill_value self.fill_leading = fill_leading - self._last_values: BufferList[Tensor] = BufferList() - self._last_indices: BufferList[Tensor] = BufferList() + self.register_buffer("_last_values", None) + self.register_buffer("_last_indices", None) def _fit_ensemble( self, @@ -51,8 +53,18 @@ def _fit_ensemble( last_values.append(last_value) last_indices.append(last_index) - self._last_values = BufferList(last_values) - self._last_indices = BufferList(last_indices) + self._set_buffer( + "_last_values", + EnsembleTensor( + groups=last_values, locations=ensemble_table._locations + ), + ) + self._set_buffer( + "_last_indices", + EnsembleTensor( + groups=last_indices, locations=ensemble_table._locations + ), + ) def _transform_ensemble( self, @@ -73,8 +85,8 @@ def _transform_ensemble( ) x = group.numerical - last_value = self._last_values[group_id] - last_index = self._last_indices[group_id] + last_value = self._last_values._groups[group_id] + last_index = self._last_indices._groups[group_id] offset = self._offset[first_member_id] # Prepend the context anchor only when processing future rows. diff --git a/sdm/processing/numerical/random_projection.py b/sdm/processing/numerical/random_projection.py index e26e2f766..f315a6cd0 100644 --- a/sdm/processing/numerical/random_projection.py +++ b/sdm/processing/numerical/random_projection.py @@ -5,8 +5,7 @@ import torch -from sdm import EnsembleTable, Stype, TableTensor -from sdm.nn._buffer import BufferList +from sdm import EnsembleTable, EnsembleTensor, Stype, TableTensor from sdm.processing import EnsembleProcessor @@ -23,6 +22,8 @@ class RandomProjection(EnsembleProcessor): handles_stypes = (Stype.numerical,) requires_fit = True + _weights: EnsembleTensor[torch.Tensor] + def __init__( self, channels: int, @@ -32,7 +33,7 @@ def __init__( super().__init__() self.channels = channels self.init = init - self._weights: BufferList[torch.Tensor] = BufferList() + self.register_buffer("_weights", None) def _fit_ensemble( self, @@ -54,20 +55,27 @@ def _fit_ensemble( assert self.init == "normal" weight.normal_(std=self.channels**-0.5, generator=generator) weights.append(weight) - self._weights = BufferList(weights) + positions = [0] * ensemble_table.num_groups + locations = [] + for group_id, _ in ensemble_table._locations: + locations.append((group_id, positions[group_id])) + positions[group_id] += 1 + self._set_buffer( + "_weights", EnsembleTensor(groups=weights, locations=locations) + ) def _transform_ensemble( self, ensemble_table: EnsembleTable, ) -> EnsembleTable: groups = [] - for i in range(ensemble_table.num_groups): + for i, weight in enumerate(self._weights._iter_groups()): group = ensemble_table.expanded_group(i) projected = TableTensor( columns={ Stype.numerical: [f"rp_{i}" for i in range(self.channels)] }, - numerical=group.numerical @ self._weights[i].transpose(-1, -2), + numerical=group.numerical @ weight.transpose(-1, -2), ) group = cast( TableTensor, diff --git a/test/processing/test_ensemble.py b/test/processing/test_ensemble.py index b8d555bfc..58e7925b5 100644 --- a/test/processing/test_ensemble.py +++ b/test/processing/test_ensemble.py @@ -1,10 +1,13 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import copy +from typing import cast + import pytest import torch -from sdm import EnsembleTable, Stype, TableTensor +from sdm import EnsembleTable, EnsembleTensor, Stype, TableTensor from sdm.processing import ( EnsembleInvertibleMixin, EnsembleProcessor, @@ -13,6 +16,7 @@ Processor, Standardize, ) +from sdm.testing import withCUDA # TODO: Replace these stubs with real EnsembleProcessor subclasses once they @@ -75,6 +79,56 @@ def _inverse_transform(self, table: TableTensor) -> TableTensor: return self._transform(table) +def test_ensemble_buffer_replacement_and_loading() -> None: + processor = IdentityEnsembleProcessor() + processor.register_buffer("values", None) + processor.values = torch.zeros(1) + restored = IdentityEnsembleProcessor() + restored.register_buffer("values", None) + restored.load_state_dict(processor.state_dict()) + torch.testing.assert_close(restored.values, torch.zeros(1)) + processor.register_buffer( + "values", + EnsembleTensor( + groups=(torch.arange(6.0).reshape(2, 3),), + locations=((0, 0), (0, 1)), + ), + ) + saved = copy.deepcopy(processor.state_dict()) + processor.values = EnsembleTensor( + groups=(torch.ones(1, 1),), locations=((0, 0),) + ) + processor.load_state_dict(saved) + actual = cast(EnsembleTensor[torch.Tensor], processor.values) + assert isinstance(actual, EnsembleTensor) + torch.testing.assert_close(actual[1], torch.tensor([3.0, 4.0, 5.0])) + cloned = copy.deepcopy(processor) + cloned_values = cast(EnsembleTensor[torch.Tensor], cloned.values) + assert isinstance(cloned_values, EnsembleTensor) + cloned_values[0].zero_() + torch.testing.assert_close(actual[0], torch.tensor([0.0, 1.0, 2.0])) + processor.register_buffer("values", torch.ones(1)) + assert isinstance(processor.values, torch.Tensor) + + +@withCUDA +def test_nonpersistent_ensemble_buffer(device: torch.device) -> None: + processor = IdentityEnsembleProcessor() + processor.register_buffer( + "temporary", + EnsembleTensor(groups=(torch.ones(1, 2),), locations=((0, 0),)), + persistent=False, + ) + processor.to(device) + temporary = cast(EnsembleTensor[torch.Tensor], processor.temporary) + assert isinstance(temporary, EnsembleTensor) + assert temporary[0].device == device + assert not any( + key.startswith("temporary.") for key in processor.state_dict() + ) + processor.load_state_dict(processor.state_dict()) + + def test_ensemble_processor_preserves_member_order_and_metadata() -> None: first = TableTensor.from_tensor( torch.tensor([[1.0], [2.0]]),