From 1b5b183162fb40f6b725463f8d80d21180260402 Mon Sep 17 00:00:00 2001 From: RBendias Date: Thu, 5 Nov 2026 13:59:17 +0000 Subject: [PATCH 1/4] Support registered ensemble tensor buffers in processors --- sdm/nn/_buffer.py | 92 ++++++++++++++++++- sdm/processing/common/shuffle.py | 16 +++- sdm/processing/ensemble.py | 69 +++++++++++++- sdm/processing/numerical/detrend.py | 37 ++++++-- sdm/processing/numerical/flip_sign.py | 14 ++- sdm/processing/numerical/interpolate.py | 31 +++++-- sdm/processing/numerical/random_projection.py | 21 +++-- test/processing/test_ensemble.py | 81 +++++++++++++++- 8 files changed, 325 insertions(+), 36 deletions(-) diff --git a/sdm/nn/_buffer.py b/sdm/nn/_buffer.py index 38e07bc6c5..39dd90008b 100644 --- a/sdm/nn/_buffer.py +++ b/sdm/nn/_buffer.py @@ -3,15 +3,103 @@ 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]") +class _EnsembleBuffer(torch.nn.Module, EnsembleTensor[Tensor]): + 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() + ) + + class BufferList(torch.nn.Module, Generic[T]): """Ordered list of tensors registered as PyTorch module state. diff --git a/sdm/processing/common/shuffle.py b/sdm/processing/common/shuffle.py index 856641ef4a..d224ac8a30 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 @@ -32,7 +31,10 @@ def __init__( ) -> None: super().__init__() self.method = method - self._permutations: BufferList[Tensor] = BufferList() + self._permutations: EnsembleTensor[Tensor] + self.register_buffer( + "_permutations", EnsembleTensor(groups=(), locations=()) + ) self._host_permutations: tuple[tuple[int, ...], ...] = () def get_extra_state(self) -> tuple[tuple[int, ...], ...]: @@ -92,7 +94,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 9ece89286e..56e4698ff9 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 _EnsembleBuffer from sdm.processing import InvertibleMixin, Processor @@ -31,6 +33,69 @@ 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, + _EnsembleBuffer( + 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), _EnsembleBuffer): + 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, _EnsembleBuffer) 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, + ) + @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 a3d8eaaacf..d9901d9794 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 @@ -25,8 +24,14 @@ class Detrend(SeriesEnsembleProcessor, EnsembleInvertibleMixin): def __init__(self, threshold: float | None = None) -> None: super().__init__() self.threshold = threshold - self._slopes: BufferList[Tensor] = BufferList() - self._intercepts: BufferList[Tensor] = BufferList() + self._slopes: EnsembleTensor[Tensor] + self._intercepts: EnsembleTensor[Tensor] + self.register_buffer( + "_slopes", EnsembleTensor(groups=(), locations=()) + ) + self.register_buffer( + "_intercepts", EnsembleTensor(groups=(), locations=()) + ) def _fit_ensemble( self, @@ -66,8 +71,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 +103,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 +133,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 57cec8958c..0a2071a17a 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 @@ -25,7 +24,8 @@ def __init__( ) -> None: super().__init__() self.probability = probability - self._signs: BufferList[Tensor] = BufferList() + self._signs: EnsembleTensor[Tensor] + self.register_buffer("_signs", EnsembleTensor(groups=(), locations=())) def _fit_ensemble( self, @@ -42,7 +42,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 540f5ed1ab..19343598ec 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 @@ -31,8 +30,14 @@ 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._last_values: EnsembleTensor[Tensor] + self._last_indices: EnsembleTensor[Tensor] + self.register_buffer( + "_last_values", EnsembleTensor(groups=(), locations=()) + ) + self.register_buffer( + "_last_indices", EnsembleTensor(groups=(), locations=()) + ) def _fit_ensemble( self, @@ -51,8 +56,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 +88,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 2ad600422a..c9e1062167 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 @@ -32,7 +31,10 @@ def __init__( super().__init__() self.channels = channels self.init = init - self._weights: BufferList[torch.Tensor] = BufferList() + self._weights: EnsembleTensor[torch.Tensor] + self.register_buffer( + "_weights", EnsembleTensor(groups=(), locations=()) + ) def _fit_ensemble( self, @@ -54,20 +56,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 b8d555bfc5..7e8265d0e7 100644 --- a/test/processing/test_ensemble.py +++ b/test/processing/test_ensemble.py @@ -1,10 +1,14 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import copy +import io +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 +17,7 @@ Processor, Standardize, ) +from sdm.testing import withCUDA # TODO: Replace these stubs with real EnsembleProcessor subclasses once they @@ -75,6 +80,80 @@ def _inverse_transform(self, table: TableTensor) -> TableTensor: return self._transform(table) +@withCUDA +def test_ensemble_buffer_roundtrip(device: torch.device) -> None: + values = EnsembleTensor( + groups=(torch.tensor([[1.0], [2.0]]), torch.tensor([[3.0, 4.0]])), + locations=((0, 1), (1, 0), (0, 0), (0, 1)), + ) + source = IdentityEnsembleProcessor() + source.register_buffer("values", values) + source.to(device=device, dtype=torch.float64) + saved = io.BytesIO() + torch.save(torch.nn.Sequential(source).state_dict(), saved) + saved.seek(0) + restored = IdentityEnsembleProcessor() + restored.register_buffer("values", EnsembleTensor(groups=(), locations=())) + torch.nn.Sequential(restored).load_state_dict( + torch.load(saved, weights_only=True, map_location="cpu") + ) + restored.to(device) + actual = cast(EnsembleTensor[torch.Tensor], restored.values) + assert isinstance(actual, EnsembleTensor) + assert len(actual) == len(values) + for index in range(len(values)): + torch.testing.assert_close( + actual[index], values[index].to(device).double() + ) + assert actual[0].data_ptr() == actual[3].data_ptr() + assert all(buffer.device == device for buffer in restored.buffers()) + + +def test_ensemble_buffer_replacement_and_loading() -> None: + processor = IdentityEnsembleProcessor() + processor.register_buffer("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]]), From c1decf3ac342e8086504beec8e3d906766a95751 Mon Sep 17 00:00:00 2001 From: RBendias Date: Thu, 5 Nov 2026 19:27:54 +0000 Subject: [PATCH 2/4] Defer ensemble buffer initialization until fitting --- sdm/nn/_buffer.py | 2 ++ sdm/processing/common/shuffle.py | 4 +-- sdm/processing/ensemble.py | 33 ++++++++++++++++++- sdm/processing/numerical/detrend.py | 8 ++--- sdm/processing/numerical/flip_sign.py | 2 +- sdm/processing/numerical/interpolate.py | 8 ++--- sdm/processing/numerical/random_projection.py | 4 +-- test/processing/test_ensemble.py | 9 +++-- 8 files changed, 48 insertions(+), 22 deletions(-) diff --git a/sdm/nn/_buffer.py b/sdm/nn/_buffer.py index 39dd90008b..67e48987b3 100644 --- a/sdm/nn/_buffer.py +++ b/sdm/nn/_buffer.py @@ -15,6 +15,8 @@ class _EnsembleBuffer(torch.nn.Module, EnsembleTensor[Tensor]): + """Adapt ensemble tensor storage to PyTorch's module lifecycle.""" + def __init__( self, groups: Sequence[Tensor], diff --git a/sdm/processing/common/shuffle.py b/sdm/processing/common/shuffle.py index d224ac8a30..630cbf33d4 100644 --- a/sdm/processing/common/shuffle.py +++ b/sdm/processing/common/shuffle.py @@ -32,9 +32,7 @@ def __init__( super().__init__() self.method = method self._permutations: EnsembleTensor[Tensor] - self.register_buffer( - "_permutations", EnsembleTensor(groups=(), locations=()) - ) + self.register_buffer("_permutations", None) self._host_permutations: tuple[tuple[int, ...], ...] = () def get_extra_state(self) -> tuple[tuple[int, ...], ...]: diff --git a/sdm/processing/ensemble.py b/sdm/processing/ensemble.py index 56e4698ff9..445e0cef55 100644 --- a/sdm/processing/ensemble.py +++ b/sdm/processing/ensemble.py @@ -48,7 +48,7 @@ def register_buffer( persistent: Whether to include the buffer in ``state_dict``. """ if isinstance(tensor, EnsembleTensor): - tensor = cast(EnsembleTensor[Tensor], tensor) + tensor = cast("EnsembleTensor[Tensor]", tensor) self._buffers.pop(name, None) self.add_module( name, @@ -96,6 +96,37 @@ def _set_buffer( 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 d9901d9794..34baf30623 100644 --- a/sdm/processing/numerical/detrend.py +++ b/sdm/processing/numerical/detrend.py @@ -26,12 +26,8 @@ def __init__(self, threshold: float | None = None) -> None: self.threshold = threshold self._slopes: EnsembleTensor[Tensor] self._intercepts: EnsembleTensor[Tensor] - self.register_buffer( - "_slopes", EnsembleTensor(groups=(), locations=()) - ) - self.register_buffer( - "_intercepts", EnsembleTensor(groups=(), locations=()) - ) + self.register_buffer("_slopes", None) + self.register_buffer("_intercepts", None) def _fit_ensemble( self, diff --git a/sdm/processing/numerical/flip_sign.py b/sdm/processing/numerical/flip_sign.py index 0a2071a17a..72c3fc782b 100644 --- a/sdm/processing/numerical/flip_sign.py +++ b/sdm/processing/numerical/flip_sign.py @@ -25,7 +25,7 @@ def __init__( super().__init__() self.probability = probability self._signs: EnsembleTensor[Tensor] - self.register_buffer("_signs", EnsembleTensor(groups=(), locations=())) + self.register_buffer("_signs", None) def _fit_ensemble( self, diff --git a/sdm/processing/numerical/interpolate.py b/sdm/processing/numerical/interpolate.py index 19343598ec..67d4ccb849 100644 --- a/sdm/processing/numerical/interpolate.py +++ b/sdm/processing/numerical/interpolate.py @@ -32,12 +32,8 @@ def __init__( self.fill_leading = fill_leading self._last_values: EnsembleTensor[Tensor] self._last_indices: EnsembleTensor[Tensor] - self.register_buffer( - "_last_values", EnsembleTensor(groups=(), locations=()) - ) - self.register_buffer( - "_last_indices", EnsembleTensor(groups=(), locations=()) - ) + self.register_buffer("_last_values", None) + self.register_buffer("_last_indices", None) def _fit_ensemble( self, diff --git a/sdm/processing/numerical/random_projection.py b/sdm/processing/numerical/random_projection.py index c9e1062167..f1f2d98af5 100644 --- a/sdm/processing/numerical/random_projection.py +++ b/sdm/processing/numerical/random_projection.py @@ -32,9 +32,7 @@ def __init__( self.channels = channels self.init = init self._weights: EnsembleTensor[torch.Tensor] - self.register_buffer( - "_weights", EnsembleTensor(groups=(), locations=()) - ) + self.register_buffer("_weights", None) def _fit_ensemble( self, diff --git a/test/processing/test_ensemble.py b/test/processing/test_ensemble.py index 7e8265d0e7..99a759aa9c 100644 --- a/test/processing/test_ensemble.py +++ b/test/processing/test_ensemble.py @@ -93,7 +93,7 @@ def test_ensemble_buffer_roundtrip(device: torch.device) -> None: torch.save(torch.nn.Sequential(source).state_dict(), saved) saved.seek(0) restored = IdentityEnsembleProcessor() - restored.register_buffer("values", EnsembleTensor(groups=(), locations=())) + restored.register_buffer("values", None) torch.nn.Sequential(restored).load_state_dict( torch.load(saved, weights_only=True, map_location="cpu") ) @@ -111,7 +111,12 @@ def test_ensemble_buffer_roundtrip(device: torch.device) -> None: def test_ensemble_buffer_replacement_and_loading() -> None: processor = IdentityEnsembleProcessor() - processor.register_buffer("values", torch.zeros(1)) + 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( From d8b4e73ac1e1bb3c36b7224f4e8500e57a59f470 Mon Sep 17 00:00:00 2001 From: RBendias Date: Fri, 6 Nov 2026 12:29:17 +0000 Subject: [PATCH 3/4] Clarify ensemble tensor buffer naming and simplify coverage --- sdm/nn/_buffer.py | 2 +- sdm/processing/common/shuffle.py | 3 +- sdm/processing/ensemble.py | 8 ++--- sdm/processing/numerical/detrend.py | 5 ++-- sdm/processing/numerical/flip_sign.py | 3 +- sdm/processing/numerical/interpolate.py | 5 ++-- sdm/processing/numerical/random_projection.py | 3 +- test/processing/test_ensemble.py | 30 ------------------- 8 files changed, 17 insertions(+), 42 deletions(-) diff --git a/sdm/nn/_buffer.py b/sdm/nn/_buffer.py index 67e48987b3..6facd6fc21 100644 --- a/sdm/nn/_buffer.py +++ b/sdm/nn/_buffer.py @@ -14,7 +14,7 @@ T = TypeVar("T", bound="Tensor | BufferList[Any]") -class _EnsembleBuffer(torch.nn.Module, EnsembleTensor[Tensor]): +class _EnsembleTensorBuffer(torch.nn.Module, EnsembleTensor[Tensor]): """Adapt ensemble tensor storage to PyTorch's module lifecycle.""" def __init__( diff --git a/sdm/processing/common/shuffle.py b/sdm/processing/common/shuffle.py index 630cbf33d4..9a66172a70 100644 --- a/sdm/processing/common/shuffle.py +++ b/sdm/processing/common/shuffle.py @@ -25,13 +25,14 @@ 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: EnsembleTensor[Tensor] self.register_buffer("_permutations", None) self._host_permutations: tuple[tuple[int, ...], ...] = () diff --git a/sdm/processing/ensemble.py b/sdm/processing/ensemble.py index 445e0cef55..316403dd55 100644 --- a/sdm/processing/ensemble.py +++ b/sdm/processing/ensemble.py @@ -11,7 +11,7 @@ from torch import Tensor from sdm import EnsembleTable, EnsembleTensor, Stype, TableTensor -from sdm.nn._buffer import _EnsembleBuffer +from sdm.nn._buffer import _EnsembleTensorBuffer from sdm.processing import InvertibleMixin, Processor @@ -52,7 +52,7 @@ def register_buffer( self._buffers.pop(name, None) self.add_module( name, - _EnsembleBuffer( + _EnsembleTensorBuffer( groups=tensor._groups, locations=tensor._locations, persistent=persistent, @@ -63,13 +63,13 @@ def register_buffer( else: self._non_persistent_buffers_set.add(name) return - if isinstance(self._modules.get(name), _EnsembleBuffer): + 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, _EnsembleBuffer) and ( + if isinstance(module, _EnsembleTensorBuffer) and ( value is None or isinstance(value, (Tensor, EnsembleTensor)) ): self.register_buffer(name, value, persistent=module._persistent) diff --git a/sdm/processing/numerical/detrend.py b/sdm/processing/numerical/detrend.py index 34baf30623..0f9daf097c 100644 --- a/sdm/processing/numerical/detrend.py +++ b/sdm/processing/numerical/detrend.py @@ -21,11 +21,12 @@ 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: EnsembleTensor[Tensor] - self._intercepts: EnsembleTensor[Tensor] self.register_buffer("_slopes", None) self.register_buffer("_intercepts", None) diff --git a/sdm/processing/numerical/flip_sign.py b/sdm/processing/numerical/flip_sign.py index 72c3fc782b..30fb823afb 100644 --- a/sdm/processing/numerical/flip_sign.py +++ b/sdm/processing/numerical/flip_sign.py @@ -18,13 +18,14 @@ 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: EnsembleTensor[Tensor] self.register_buffer("_signs", None) def _fit_ensemble( diff --git a/sdm/processing/numerical/interpolate.py b/sdm/processing/numerical/interpolate.py index 67d4ccb849..ded90301f6 100644 --- a/sdm/processing/numerical/interpolate.py +++ b/sdm/processing/numerical/interpolate.py @@ -21,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, *, @@ -30,8 +33,6 @@ def __init__( super().__init__() self.fill_value = fill_value self.fill_leading = fill_leading - self._last_values: EnsembleTensor[Tensor] - self._last_indices: EnsembleTensor[Tensor] self.register_buffer("_last_values", None) self.register_buffer("_last_indices", None) diff --git a/sdm/processing/numerical/random_projection.py b/sdm/processing/numerical/random_projection.py index f1f2d98af5..487fb619a1 100644 --- a/sdm/processing/numerical/random_projection.py +++ b/sdm/processing/numerical/random_projection.py @@ -22,6 +22,8 @@ class RandomProjection(EnsembleProcessor): handles_stypes = (Stype.numerical,) requires_fit = True + _weights: EnsembleTensor[torch.Tensor] + def __init__( self, channels: int, @@ -31,7 +33,6 @@ def __init__( super().__init__() self.channels = channels self.init = init - self._weights: EnsembleTensor[torch.Tensor] self.register_buffer("_weights", None) def _fit_ensemble( diff --git a/test/processing/test_ensemble.py b/test/processing/test_ensemble.py index 99a759aa9c..58e7925b5b 100644 --- a/test/processing/test_ensemble.py +++ b/test/processing/test_ensemble.py @@ -2,7 +2,6 @@ # SPDX-License-Identifier: Apache-2.0 import copy -import io from typing import cast import pytest @@ -80,35 +79,6 @@ def _inverse_transform(self, table: TableTensor) -> TableTensor: return self._transform(table) -@withCUDA -def test_ensemble_buffer_roundtrip(device: torch.device) -> None: - values = EnsembleTensor( - groups=(torch.tensor([[1.0], [2.0]]), torch.tensor([[3.0, 4.0]])), - locations=((0, 1), (1, 0), (0, 0), (0, 1)), - ) - source = IdentityEnsembleProcessor() - source.register_buffer("values", values) - source.to(device=device, dtype=torch.float64) - saved = io.BytesIO() - torch.save(torch.nn.Sequential(source).state_dict(), saved) - saved.seek(0) - restored = IdentityEnsembleProcessor() - restored.register_buffer("values", None) - torch.nn.Sequential(restored).load_state_dict( - torch.load(saved, weights_only=True, map_location="cpu") - ) - restored.to(device) - actual = cast(EnsembleTensor[torch.Tensor], restored.values) - assert isinstance(actual, EnsembleTensor) - assert len(actual) == len(values) - for index in range(len(values)): - torch.testing.assert_close( - actual[index], values[index].to(device).double() - ) - assert actual[0].data_ptr() == actual[3].data_ptr() - assert all(buffer.device == device for buffer in restored.buffers()) - - def test_ensemble_buffer_replacement_and_loading() -> None: processor = IdentityEnsembleProcessor() processor.register_buffer("values", None) From 76909cd2e44f68dc13d303e4d6bf1642aed0e30c Mon Sep 17 00:00:00 2001 From: RBendias Date: Fri, 6 Nov 2026 12:59:55 +0000 Subject: [PATCH 4/4] Present public buffer API first and clarify ensemble buffer docs --- sdm/nn/_buffer.py | 160 +++++++++++++++++++++++----------------------- 1 file changed, 80 insertions(+), 80 deletions(-) diff --git a/sdm/nn/_buffer.py b/sdm/nn/_buffer.py index 6facd6fc21..c0ca70b5d4 100644 --- a/sdm/nn/_buffer.py +++ b/sdm/nn/_buffer.py @@ -14,8 +14,87 @@ T = TypeVar("T", bound="Tensor | BufferList[Any]") +class BufferList(torch.nn.Module, Generic[T]): + """Ordered list of tensors registered as PyTorch module state. + + Args: + buffers: Tensors and buffer lists to register in order. + """ + + def __init__( + self, + buffers: Iterable[T] = (), + ) -> None: + super().__init__() + for index, item in enumerate(buffers): + if isinstance(item, BufferList): + self.add_module(str(index), item) + else: + self.register_buffer(str(index), item) + + 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: + layout = state_dict[f"{prefix}_extra_state"] + materialized_keys = [] + for index, is_nested in enumerate(layout): + if is_nested: + continue + name = str(index) + key = f"{prefix}{index}" + if name not in self._buffers: + # Some tensor subclasses cannot use empty_like or copy_. + self.register_buffer(name, state_dict.pop(key).clone()) + materialized_keys.append(key) + + super()._load_from_state_dict( + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ) + for key in materialized_keys: + if key in missing_keys: + missing_keys.remove(key) + + def __getitem__(self, index: int) -> T: + name = str(index) + if name in self._modules: + return cast(T, self.get_submodule(name)) + return cast(T, self.get_buffer(name)) + + def __iter__(self) -> Iterator[T]: + return (self[index] for index in range(len(self))) + + def __len__(self) -> int: + return len(self._buffers) + len(self._modules) + + def get_extra_state(self) -> tuple[bool, ...]: + r""":meta private:""" # noqa: D415 + # Layout of the buffer list in a tuple of booleans. + # True at positions that hold a nested BufferList rather than a tensor. + return tuple(str(index) in self._modules for index in range(len(self))) + + def set_extra_state(self, state: object) -> None: + r""":meta private:""" # noqa: D415 + for index, is_nested in enumerate(cast(tuple[bool, ...], state)): + name = str(index) + if is_nested and name not in self._modules: + self.add_module(name, BufferList()) + + class _EnsembleTensorBuffer(torch.nn.Module, EnsembleTensor[Tensor]): - """Adapt ensemble tensor storage to PyTorch's module lifecycle.""" + """Ensemble tensor registered as PyTorch module state.""" def __init__( self, @@ -100,82 +179,3 @@ def _load_from_state_dict( self._groups = tuple( cast(Tensor, group) for group in self._buffers.values() ) - - -class BufferList(torch.nn.Module, Generic[T]): - """Ordered list of tensors registered as PyTorch module state. - - Args: - buffers: Tensors and buffer lists to register in order. - """ - - def __init__( - self, - buffers: Iterable[T] = (), - ) -> None: - super().__init__() - for index, item in enumerate(buffers): - if isinstance(item, BufferList): - self.add_module(str(index), item) - else: - self.register_buffer(str(index), item) - - 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: - layout = state_dict[f"{prefix}_extra_state"] - materialized_keys = [] - for index, is_nested in enumerate(layout): - if is_nested: - continue - name = str(index) - key = f"{prefix}{index}" - if name not in self._buffers: - # Some tensor subclasses cannot use empty_like or copy_. - self.register_buffer(name, state_dict.pop(key).clone()) - materialized_keys.append(key) - - super()._load_from_state_dict( - state_dict, - prefix, - local_metadata, - strict, - missing_keys, - unexpected_keys, - error_msgs, - ) - for key in materialized_keys: - if key in missing_keys: - missing_keys.remove(key) - - def __getitem__(self, index: int) -> T: - name = str(index) - if name in self._modules: - return cast(T, self.get_submodule(name)) - return cast(T, self.get_buffer(name)) - - def __iter__(self) -> Iterator[T]: - return (self[index] for index in range(len(self))) - - def __len__(self) -> int: - return len(self._buffers) + len(self._modules) - - def get_extra_state(self) -> tuple[bool, ...]: - r""":meta private:""" # noqa: D415 - # Layout of the buffer list in a tuple of booleans. - # True at positions that hold a nested BufferList rather than a tensor. - return tuple(str(index) in self._modules for index in range(len(self))) - - def set_extra_state(self, state: object) -> None: - r""":meta private:""" # noqa: D415 - for index, is_nested in enumerate(cast(tuple[bool, ...], state)): - name = str(index) - if is_nested and name not in self._modules: - self.add_module(name, BufferList())