Fix processor tracing and fitted-buffer updates

sdm/processing/base.py

@@ -32,7 +32,7 @@ class Processor(torch.nn.Module, abc.ABC):
     """
 
     #: Semantic types this processor operates on.
-    handles_stypes: frozenset[Stype]
+    handles_stypes: tuple[Stype, ...]
 
     #: Whether this processor requires fitting.
     requires_fit: bool
@@ -161,7 +161,12 @@ def _fit_transform(
         return self._transform(table)
 
     def _handles_table(self, table: TableTensor) -> bool:
-        return bool(table.active_stypes & self.handles_stypes)
+        # Check column counts directly so torch.compile can trace this without
+        # constructing or intersecting sets of column types.
+        for stype in self.handles_stypes:
+            if getattr(table, stype.value).size(-1) > 0:
+                return True
+        return False
 
     def fit(
         self,

sdm/processing/categorical/align.py

@@ -68,7 +68,7 @@ class AlignCategories(EnsembleProcessor):
     ``"green"`` and the missing value use ``-1``.
     """
 
-    handles_stypes = frozenset({Stype.categorical})
+    handles_stypes = (Stype.categorical,)
     requires_fit = True
 
     def __init__(

sdm/processing/categorical/count.py

@@ -32,9 +32,12 @@ class AddCategoryCounts(Processor):
             vocabulary size. Defaults to 0.
     """
 
-    handles_stypes = frozenset({Stype.categorical})
+    handles_stypes = (Stype.categorical,)
     requires_fit = True
 
+    _log_counts: Tensor
+    _num_categories: Tensor
+
     def __init__(self, min_cardinality: int = 0) -> None:
         super().__init__()
         self.min_cardinality = min_cardinality
@@ -89,8 +92,11 @@ def _fit(
             if dtype in {torch.float16, torch.bfloat16}
             else dtype
         )
-        self._log_counts = counts.to(count_dtype).log1p().to(dtype)
-        self._num_categories = num_categories
+        self._set_buffer(
+            name="_log_counts",
+            value=counts.to(count_dtype).log1p().to(dtype),
+        )
+        self._set_buffer("_num_categories", num_categories)
         self._categories = BufferList(table.categorical.categories)
 
     def _transform(self, table: TableTensor) -> TableTensor:

sdm/processing/categorical/impute.py

@@ -25,9 +25,11 @@ class ImputeMode(Processor):
     training and transform inputs were tensorized independently.
     """
 
-    handles_stypes = frozenset({Stype.categorical})
+    handles_stypes = (Stype.categorical,)
     requires_fit = True
 
+    _fill_values: torch.Tensor
+
     def __init__(self) -> None:
         super().__init__()
         # TODO: Add a separate processor that encodes missing values as their
@@ -71,10 +73,11 @@ def _fit(
             )
             fill_values.append(counts.argmax(dim=-1, keepdim=True))
 
-        self._fill_values = (
-            torch.stack(fill_values, dim=-1)
+        self._set_buffer(
+            name="_fill_values",
+            value=torch.stack(fill_values, dim=-1)
             if len(fill_values) > 0
-            else torch.empty(0, dtype=torch.long, device=data.device)
+            else torch.empty(0, dtype=torch.long, device=data.device),
         )
         self._categories = BufferList(table.categorical.categories)
 

sdm/processing/categorical/shuffle.py

@@ -27,7 +27,7 @@ class ShuffleCategories(EnsembleProcessor):
             offset, cycling through every offset in random order.
     """
 
-    handles_stypes = frozenset({Stype.categorical})
+    handles_stypes = (Stype.categorical,)
     requires_fit = True
 
     def __init__(

sdm/processing/common/callable.py

@@ -10,7 +10,7 @@
 class Callable(Processor):
     """Adapt a stateless callable to the :class:`Processor` interface."""
 
-    handles_stypes = frozenset(Stype)
+    handles_stypes = tuple(Stype)
     requires_fit = False
 
     def __init__(self, function: Callable[[TableTensor], TableTensor]) -> None:

sdm/processing/common/choice.py

@@ -36,9 +36,7 @@ def __init__(
         for arg in args:
             options.append(EnsembleProcessor.as_processor(arg))
         self.options: ModuleList[EnsembleProcessor] = ModuleList(options)
-        self.handles_stypes = frozenset(
-            stype for option in self.options for stype in option.handles_stypes
-        )
+        self.handles_stypes = self._merge_stypes(self.options)
         self.method = method
         self._option_ids: tuple[int, ...] = ()
 

sdm/processing/common/drop.py

@@ -18,7 +18,9 @@ class DropStypes(Processor):
     def __init__(self, *stypes: StypeLike) -> None:
         super().__init__()
         self._stypes = stypes
-        self.handles_stypes = frozenset(Stype(stype) for stype in stypes)
+        self.handles_stypes = tuple(
+            dict.fromkeys(Stype(stype) for stype in stypes)
+        )
 
     def _transform(self, table: TableTensor) -> TableTensor:
         return table.drop_stypes(self.handles_stypes)

sdm/processing/common/ensemble.py

@@ -60,7 +60,7 @@ def set_extra_state(self, state: object) -> None:
         )
 
     @property
-    def handles_stypes(self) -> frozenset[Stype]:
+    def handles_stypes(self) -> tuple[Stype, ...]:
         r""":meta private:"""  # noqa: D415
         return self.processor.handles_stypes
 

sdm/processing/common/identity.py

@@ -8,7 +8,7 @@
 class Identity(Processor, InvertibleMixin):
     r"""Return inputs unchanged."""
 
-    handles_stypes = frozenset(Stype)
+    handles_stypes = tuple(Stype)
     requires_fit = False
 
     def _transform(self, table: TableTensor) -> TableTensor:

sdm/processing/common/select.py

@@ -23,7 +23,7 @@ class SelectColumns(EnsembleProcessor):
             0, 1, and 2 receive columns (0, 1), (2, 3), and (4, 0).
     """
 
-    handles_stypes = frozenset(Stype)
+    handles_stypes = tuple(Stype)
     requires_fit = False
 
     def __init__(

sdm/processing/common/sequential.py

@@ -26,13 +26,11 @@ def __init__(self, *args: object) -> None:
         self.extend(args)
 
     @property
-    def handles_stypes(self) -> frozenset[Stype]:
+    def handles_stypes(self) -> tuple[Stype, ...]:
         r""":meta private:"""  # noqa: D415
         if len(self) == 0:
-            return frozenset(Stype)
-        return frozenset(
-            stype for child in self for stype in child.handles_stypes
-        )
+            return tuple(Stype)
+        return self._merge_stypes(self)
 
     def append(self, processor: object) -> Self:
         r"""Append a processor or callable to this sequence.

sdm/processing/common/shuffle.py

@@ -23,7 +23,7 @@ class ShuffleColumns(EnsembleProcessor, EnsembleInvertibleMixin):
             permutations.
     """
 
-    handles_stypes = frozenset({Stype.numerical})
+    handles_stypes = (Stype.numerical,)
     requires_fit = True
 
     def __init__(

sdm/processing/common/stype.py

@@ -64,9 +64,7 @@ def __init__(
             processor = EnsembleProcessor.as_processor(processor)
             self.processors[str(stype)] = processor
 
-        self.handles_stypes = frozenset(
-            Stype(stype) for stype in self.processors
-        )
+        self.handles_stypes = tuple(Stype(stype) for stype in self.processors)
         self.requires_fit = any(
             processor.requires_fit for processor in self.processors.values()
         )

sdm/processing/common/table.py

@@ -39,13 +39,9 @@ def __init__(
         self._route: Literal["task", "related"] | None = None
 
     @property
-    def handles_stypes(self) -> frozenset[Stype]:
+    def handles_stypes(self) -> tuple[Stype, ...]:
         r""":meta private:"""  # noqa: D415
-        return frozenset(
-            stype
-            for processor in self.processors.values()
-            for stype in processor.handles_stypes
-        )
+        return self._merge_stypes(self.processors.values())
 
     def _fit_ensemble(
         self,

sdm/processing/common/task.py

@@ -43,16 +43,12 @@ def __init__(
         self._task: Literal["classification", "regression"] | None = None
 
     @property
-    def handles_stypes(self) -> frozenset[Stype]:
+    def handles_stypes(self) -> tuple[Stype, ...]:
         r""":meta private:"""  # noqa: D415
         if self._task is None:
-            return frozenset(
-                stype
-                for processor in self.processors.values()
-                for stype in processor.handles_stypes
-            )
+            return self._merge_stypes(self.processors.values())
         if self._task not in self.processors:
-            return frozenset()
+            return ()
         return self.processors[self._task].handles_stypes
 
     def _processor(self) -> EnsembleProcessor:

sdm/processing/common/to_numerical.py

@@ -24,7 +24,7 @@ class ToNumerical(Processor):
     """
 
     requires_fit = False
-    handles_stypes = frozenset({Stype.numerical, Stype.categorical})
+    handles_stypes = (Stype.numerical, Stype.categorical)
 
     def _transform(self, table: TableTensor) -> TableTensor:
         """Return ``table`` with categorical columns moved to ``numerical``."""

sdm/processing/datetime/calendar.py

@@ -24,7 +24,7 @@ class AddCalendarFields(Processor):
             ``"cyclic"`` uses sine and cosine values.
     """
 
-    handles_stypes = frozenset({Stype.datetime})
+    handles_stypes = (Stype.datetime,)
     requires_fit = False
 
     def __init__(

sdm/processing/ensemble.py

@@ -4,11 +4,12 @@
 from __future__ import annotations
 
 import abc
+from collections.abc import Iterable
 from typing import Self
 
 import torch
 
-from sdm import EnsembleTable, TableTensor
+from sdm import EnsembleTable, Stype, TableTensor
 from sdm.processing import InvertibleMixin, Processor
 
 
@@ -46,6 +47,17 @@ def as_processor(processor: object) -> EnsembleProcessor:
             return processor
         return EnsembleProcessorAdapter(processor)
 
+    @staticmethod
+    def _merge_stypes(processors: Iterable[Processor]) -> tuple[Stype, ...]:
+        """Combine column types without duplicates, in input order."""
+        return tuple(
+            dict.fromkeys(
+                stype
+                for processor in processors
+                for stype in processor.handles_stypes
+            )
+        )
+
     def _fit(
         self,
         table: TableTensor,
@@ -115,7 +127,7 @@ def fit_ensemble(
             generator: Pseudorandom number generator used for sampling.
         """
         if not any(
-            group.active_stypes & self.handles_stypes
+            self._handles_table(group)
             for group in ensemble_table._iter_groups()
         ):
             return self
@@ -137,7 +149,7 @@ def transform_ensemble(
             The transformed ensemble table.
         """
         if not any(
-            group.active_stypes & self.handles_stypes
+            self._handles_table(group)
             for group in ensemble_table._iter_groups()
         ):
             return ensemble_table
@@ -160,7 +172,7 @@ def fit_transform_ensemble(
             The transformed ensemble table.
         """
         if not any(
-            group.active_stypes & self.handles_stypes
+            self._handles_table(group)
             for group in ensemble_table._iter_groups()
         ):
             return ensemble_table

sdm/processing/execution.py

@@ -74,10 +74,13 @@ def fit_transform(
             if isinstance(module, sp.SeriesEnsembleProcessor)
         )
         for module in series_processors:
-            module._offset = torch.zeros(
-                len(y),
-                dtype=torch.long,
-                device=y[0].device,
+            module._set_buffer(
+                name="_offset",
+                value=torch.zeros(
+                    size=(len(y),),
+                    dtype=torch.long,
+                    device=y[0].device,
+                ),
             )
 
         y = self.recipe.target.fit_transform_ensemble(y, generator=generator)
@@ -138,13 +141,16 @@ def fit_transform(
                     )
                 for module in processor.modules():
                     if isinstance(module, sp.SeriesEnsembleProcessor):
-                        module._offset = torch.tensor(
-                            [
-                                ensemble_table[i].size(-2)
-                                for i in range(len(ensemble_table))
-                            ],
-                            dtype=torch.long,
-                            device=ensemble_table[0].device,
+                        module._set_buffer(
+                            name="_offset",
+                            value=torch.tensor(
+                                data=[
+                                    ensemble_table[i].size(-2)
+                                    for i in range(len(ensemble_table))
+                                ],
+                                dtype=torch.long,
+                                device=ensemble_table[0].device,
+                            ),
                         )
 
         for module in self.recipe.features.modules():
@@ -181,10 +187,13 @@ def fit_transform(
             )
 
         for module in series_processors:
-            module._offset = torch.tensor(
-                [x[i].size(-2) for i in range(len(x))],
-                dtype=torch.long,
-                device=x[0].device,
+            module._set_buffer(
+                name="_offset",
+                value=torch.tensor(
+                    data=[x[i].size(-2) for i in range(len(x))],
+                    dtype=torch.long,
+                    device=x[0].device,
+                ),
             )
 
         members: list[MemberContext] = []

sdm/processing/numerical/cast.py

@@ -14,7 +14,7 @@ class Cast(Processor):
         dtype: The floating-point dtype of the numerical columns.
     """
 
-    handles_stypes = frozenset({Stype.numerical})
+    handles_stypes = (Stype.numerical,)
     requires_fit = False
 
     def __init__(self, dtype: torch.dtype) -> None:

sdm/processing/numerical/clip.py

@@ -17,7 +17,7 @@ class Clip(Processor, InvertibleMixin):
         max_value: Inclusive upper bound for every numerical value.
     """
 
-    handles_stypes = frozenset({Stype.numerical})
+    handles_stypes = (Stype.numerical,)
     requires_fit = False
 
     def __init__(

sdm/processing/numerical/clip_soft.py

@@ -26,7 +26,7 @@ class ClipSoft(Processor):
         max_absolute_value: Bound of the soft clip.
     """
 
-    handles_stypes = frozenset({Stype.numerical})
+    handles_stypes = (Stype.numerical,)
     requires_fit = False
 
     def __init__(self, max_absolute_value: float) -> None: