From cda6da2f472c1fa20bf0536fd413ad4c91541b59 Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Thu, 15 Oct 2026 12:56:13 +0900 Subject: [PATCH 1/9] Make processor applicability checks compilable Check the column counts of handled blocks directly instead of constructing active semantic-type sets. Fit PCA, RobustScale and ImputeMean through the existing buffer registration helper to avoid tracing attribute-assignment inspection. Validated graph capture and Inductor execution on PyTorch 2.14 with static, dynamic and batched inputs, missing values and checkpoint replay. The general check still encounters unsupported frozenset iteration in PyTorch 2.7 fullgraph mode; retain Standardize's specialized check so its existing support is unchanged. Signed-off-by: Ard <88337265+aw471@users.noreply.github.com> --- sdm/processing/base.py | 6 +++++- sdm/processing/numerical/impute.py | 7 ++++++- sdm/processing/numerical/pca.py | 10 ++++++++-- sdm/processing/numerical/robust_scale.py | 8 ++++++-- 4 files changed, 25 insertions(+), 6 deletions(-) diff --git a/sdm/processing/base.py b/sdm/processing/base.py index 9c003693b9..b6fd6a53ce 100644 --- a/sdm/processing/base.py +++ b/sdm/processing/base.py @@ -161,7 +161,11 @@ def _fit_transform( return self._transform(table) def _handles_table(self, table: TableTensor) -> bool: - return bool(table.active_stypes & self.handles_stypes) + # Check handled blocks directly instead of constructing active types. + for stype in self.handles_stypes: + if getattr(table, stype.value).size(-1) > 0: + return True + return False def fit( self, diff --git a/sdm/processing/numerical/impute.py b/sdm/processing/numerical/impute.py index 319841d57c..98a9099b77 100644 --- a/sdm/processing/numerical/impute.py +++ b/sdm/processing/numerical/impute.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 import torch +from torch import Tensor from sdm import Stype, TableTensor from sdm.processing import Processor @@ -22,6 +23,8 @@ class ImputeMean(Processor): handles_stypes = frozenset({Stype.numerical}) requires_fit = True + _mean: Tensor + def __init__( self, *, @@ -43,7 +46,9 @@ def _fit( dim=-2, keepdim=True, ) - self._mean = torch.where(mean.isnan(), self.fill_value, mean) + self._set_buffer( + "_mean", torch.where(mean.isnan(), self.fill_value, mean) + ) def _transform(self, table: TableTensor) -> TableTensor: """Replace NaNs with the fitted per-column means.""" diff --git a/sdm/processing/numerical/pca.py b/sdm/processing/numerical/pca.py index 900efecfc2..e1df411d7c 100644 --- a/sdm/processing/numerical/pca.py +++ b/sdm/processing/numerical/pca.py @@ -4,6 +4,7 @@ from typing import cast import torch +from torch import Tensor from sdm import Stype, TableTensor from sdm.processing import Processor @@ -19,6 +20,9 @@ class PCA(Processor): handles_stypes = frozenset({Stype.numerical}) requires_fit = True + mean: Tensor + components: Tensor + def __init__(self, num_components: int) -> None: super().__init__() self.num_components = num_components @@ -38,7 +42,7 @@ def _fit( f"numerical features" ) - self.mean = table.numerical.mean(dim=-2, keepdim=True) + self._set_buffer("mean", table.numerical.mean(dim=-2, keepdim=True)) _, _, vh = torch.linalg.svd( table.numerical - self.mean, full_matrices=False, @@ -48,7 +52,9 @@ def _fit( table.numerical.size(-2), table.numerical.size(-1), ) - self.components = vh[..., :num_components, :].transpose(-2, -1) + self._set_buffer( + "components", vh[..., :num_components, :].transpose(-2, -1) + ) def _transform(self, table: TableTensor) -> TableTensor: x = (table.numerical - self.mean) @ self.components diff --git a/sdm/processing/numerical/robust_scale.py b/sdm/processing/numerical/robust_scale.py index 2525147db3..2043e37d95 100644 --- a/sdm/processing/numerical/robust_scale.py +++ b/sdm/processing/numerical/robust_scale.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 import torch +from torch import Tensor from sdm import Stype, TableTensor from sdm._memory import split_size @@ -24,6 +25,9 @@ class RobustScale(Processor, InvertibleMixin): handles_stypes = frozenset({Stype.numerical}) requires_fit = True + median: Tensor + scale: Tensor + def __init__( self, *, @@ -69,9 +73,9 @@ def _fit( ], dim=-1, ) - self.median = median.to(dtype=numerical.dtype) + self._set_buffer("median", median.to(dtype=numerical.dtype)) scale = torch.where(lower == upper, 1.0, upper - lower) - self.scale = scale.to(dtype=numerical.dtype) + self._set_buffer("scale", scale.to(dtype=numerical.dtype)) def _transform(self, table: TableTensor) -> TableTensor: numerical = (table.numerical - self.median).div_(self.scale) From 631bd779d73ab40809a872be3dcab65586ce26df Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Thu, 15 Oct 2026 17:30:47 +0900 Subject: [PATCH 2/9] Complete processor buffer updates and ensemble applicability checks --- sdm/processing/categorical/count.py | 9 ++++-- sdm/processing/categorical/impute.py | 7 +++-- sdm/processing/ensemble.py | 6 ++-- sdm/processing/execution.py | 35 +++++++++++++---------- sdm/processing/numerical/power.py | 33 ++++++++++++++------- sdm/processing/numerical/quantile.py | 4 +-- sdm/processing/numerical/quantile_clip.py | 8 ++++-- sdm/processing/numerical/rank_gaussian.py | 8 ++++-- sdm/processing/numerical/sigma_clip.py | 8 ++++-- sdm/processing/numerical/standardize.py | 2 +- sdm/processing/text/tfidf.py | 4 ++- 11 files changed, 81 insertions(+), 43 deletions(-) diff --git a/sdm/processing/categorical/count.py b/sdm/processing/categorical/count.py index 3a975c4600..6593a1b906 100644 --- a/sdm/processing/categorical/count.py +++ b/sdm/processing/categorical/count.py @@ -35,6 +35,9 @@ class AddCategoryCounts(Processor): handles_stypes = frozenset({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,10 @@ 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( + "_log_counts", 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: diff --git a/sdm/processing/categorical/impute.py b/sdm/processing/categorical/impute.py index 0b94fb7fd4..166ec159b2 100644 --- a/sdm/processing/categorical/impute.py +++ b/sdm/processing/categorical/impute.py @@ -28,6 +28,8 @@ class ImputeMode(Processor): handles_stypes = frozenset({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 = ( + self._set_buffer( + "_fill_values", 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) diff --git a/sdm/processing/ensemble.py b/sdm/processing/ensemble.py index 7fad1b43a1..296916f9cc 100644 --- a/sdm/processing/ensemble.py +++ b/sdm/processing/ensemble.py @@ -115,7 +115,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 +137,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 +160,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 diff --git a/sdm/processing/execution.py b/sdm/processing/execution.py index ba8f27f7d2..317d04c46a 100644 --- a/sdm/processing/execution.py +++ b/sdm/processing/execution.py @@ -74,10 +74,9 @@ 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( + "_offset", + torch.zeros(len(y), dtype=torch.long, device=y[0].device), ) y = self.recipe.target.fit_transform_ensemble(y, generator=generator) @@ -138,13 +137,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( + "_offset", + torch.tensor( + [ + 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 +183,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( + "_offset", + torch.tensor( + [x[i].size(-2) for i in range(len(x))], + dtype=torch.long, + device=x[0].device, + ), ) members: list[MemberContext] = [] diff --git a/sdm/processing/numerical/power.py b/sdm/processing/numerical/power.py index 73f34a4b2e..9c5e970e5f 100644 --- a/sdm/processing/numerical/power.py +++ b/sdm/processing/numerical/power.py @@ -229,6 +229,11 @@ class PowerTransform(Processor, InvertibleMixin): handles_stypes = frozenset({Stype.numerical}) requires_fit = True + lambdas: Tensor + max: Tensor + mean: Tensor + scale: Tensor + def __init__( self, *, @@ -272,16 +277,22 @@ def _fit( # ``mean`` never exceeds the column maximum, and is zero for an # entirely missing column, whose max must be zero. - self.max = torch.where(finite, table.numerical, mean).amax( - dim=-2, - keepdim=True, + self._set_buffer( + "max", + torch.where(finite, table.numerical, mean).amax( + dim=-2, + keepdim=True, + ), ) del finite - self.lambdas = self._optimize_lambdas( - finite_or_nan, - constant_features, - count=count, + self._set_buffer( + "lambdas", + self._optimize_lambdas( + finite_or_nan, + constant_features, + count=count, + ), ) if self.standardize: @@ -292,11 +303,11 @@ def _fit( var /= count scale = var.sqrt() scale[_constant_feature_mask(var, mean, count)] = 1.0 - self.mean = mean - self.scale = scale + self._set_buffer("mean", mean) + self._set_buffer("scale", scale) else: - self.mean = torch.zeros_like(self.lambdas) - self.scale = torch.ones_like(self.lambdas) + self._set_buffer("mean", torch.zeros_like(self.lambdas)) + self._set_buffer("scale", torch.ones_like(self.lambdas)) def _transform(self, table: TableTensor) -> TableTensor: """Transform ``table`` with fitted Yeo-Johnson parameters.""" diff --git a/sdm/processing/numerical/quantile.py b/sdm/processing/numerical/quantile.py index bc93ead525..e3f3f19bad 100644 --- a/sdm/processing/numerical/quantile.py +++ b/sdm/processing/numerical/quantile.py @@ -186,8 +186,8 @@ def _fit( ) .movedim(-1, -2) ) - self._quantiles = quantiles - self._references = references + self._set_buffer("_quantiles", quantiles) + self._set_buffer("_references", references) def _transform(self, table: TableTensor) -> TableTensor: numerical = table.numerical diff --git a/sdm/processing/numerical/quantile_clip.py b/sdm/processing/numerical/quantile_clip.py index dc425f80fe..e8305d6ea5 100644 --- a/sdm/processing/numerical/quantile_clip.py +++ b/sdm/processing/numerical/quantile_clip.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 import torch +from torch import Tensor from sdm import Stype, TableTensor from sdm.processing import Processor @@ -23,6 +24,9 @@ class ClipQuantiles(Processor): handles_stypes = frozenset({Stype.numerical}) requires_fit = True + lower_bound: Tensor + upper_bound: Tensor + def __init__( self, *, @@ -51,8 +55,8 @@ def _fit( dim=-2, keepdim=True, ) - self.lower_bound = q_low - self.upper_bound = q_high + self._set_buffer("lower_bound", q_low) + self._set_buffer("upper_bound", q_high) def _transform(self, table: TableTensor) -> TableTensor: """Clamp ``table`` to the fitted lower and upper bounds.""" diff --git a/sdm/processing/numerical/rank_gaussian.py b/sdm/processing/numerical/rank_gaussian.py index c8a0a531d4..62eef9b72d 100644 --- a/sdm/processing/numerical/rank_gaussian.py +++ b/sdm/processing/numerical/rank_gaussian.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 import torch +from torch import Tensor from sdm import Stype, TableTensor from sdm._memory import split_size @@ -34,6 +35,9 @@ class RankGaussian(Processor): handles_stypes = frozenset({Stype.numerical}) requires_fit = True + _values: Tensor + _probabilities: Tensor + def __init__(self, *, max_knots: int | None = None) -> None: super().__init__() if max_knots is not None and max_knots < 2: @@ -131,8 +135,8 @@ def _fit( ) values, probabilities = zip(*knots, strict=True) - self._values = torch.cat(values, dim=-2) - self._probabilities = torch.cat(probabilities, dim=-2) + self._set_buffer("_values", torch.cat(values, dim=-2)) + self._set_buffer("_probabilities", torch.cat(probabilities, dim=-2)) def _transform(self, table: TableTensor) -> TableTensor: numerical = table.numerical diff --git a/sdm/processing/numerical/sigma_clip.py b/sdm/processing/numerical/sigma_clip.py index 56d951a378..f477d685dd 100644 --- a/sdm/processing/numerical/sigma_clip.py +++ b/sdm/processing/numerical/sigma_clip.py @@ -4,6 +4,7 @@ import math import torch +from torch import Tensor from sdm import Stype, TableTensor from sdm._memory import split_size @@ -28,6 +29,9 @@ class ClipSigma(Processor): handles_stypes = frozenset({Stype.numerical}) requires_fit = True + lower_bound: Tensor + upper_bound: Tensor + def __init__( self, *, @@ -84,8 +88,8 @@ def _fit( mean = torch.where(has_kept, kept_mean, mean) std = torch.where(has_kept, kept_std, std) - self.lower_bound = mean - self.threshold * std - self.upper_bound = mean + self.threshold * std + self._set_buffer("lower_bound", mean - self.threshold * std) + self._set_buffer("upper_bound", mean + self.threshold * std) def _transform(self, table: TableTensor) -> TableTensor: numerical = table.numerical diff --git a/sdm/processing/numerical/standardize.py b/sdm/processing/numerical/standardize.py index f2945343ad..1fc5e3e28a 100644 --- a/sdm/processing/numerical/standardize.py +++ b/sdm/processing/numerical/standardize.py @@ -52,7 +52,7 @@ def with_state(self, mean: Tensor | float, scale: Tensor | float) -> Self: return self def _handles_table(self, table: TableTensor) -> bool: - # Check column count without constructing sets during compilation. + # PyTorch 2.7 cannot trace iteration over handled column types. return table.numerical.size(-1) > 0 def _fit( diff --git a/sdm/processing/text/tfidf.py b/sdm/processing/text/tfidf.py index 8c14f7d8a6..cde22ea392 100644 --- a/sdm/processing/text/tfidf.py +++ b/sdm/processing/text/tfidf.py @@ -49,6 +49,8 @@ class TFIDF(EnsembleProcessor): handles_stypes = frozenset({Stype.text}) requires_fit = True + _idf_template: Tensor + def __init__( self, *, @@ -310,7 +312,7 @@ def _learn_state( idfs.append(idf) if idfs: - self._idf_template = idfs[0].new_empty(0) + self._set_buffer("_idf_template", idfs[0].new_empty(0)) return _TFIDFState(vocabularies, idfs) def _fit_ensemble( From 526910b43b583acbd27e9d450c711ce2bba1e64c Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Fri, 16 Oct 2026 01:27:38 +0900 Subject: [PATCH 3/9] Use keyword arguments in multiline processor calls --- sdm/processing/categorical/count.py | 3 ++- sdm/processing/categorical/impute.py | 4 ++-- sdm/processing/execution.py | 20 ++++++++++++-------- sdm/processing/numerical/impute.py | 3 ++- sdm/processing/numerical/pca.py | 3 ++- sdm/processing/numerical/power.py | 12 ++++++------ 6 files changed, 26 insertions(+), 19 deletions(-) diff --git a/sdm/processing/categorical/count.py b/sdm/processing/categorical/count.py index 6593a1b906..1554291265 100644 --- a/sdm/processing/categorical/count.py +++ b/sdm/processing/categorical/count.py @@ -93,7 +93,8 @@ def _fit( else dtype ) self._set_buffer( - "_log_counts", counts.to(count_dtype).log1p().to(dtype) + name="_log_counts", + value=counts.to(count_dtype).log1p().to(dtype), ) self._set_buffer("_num_categories", num_categories) self._categories = BufferList(table.categorical.categories) diff --git a/sdm/processing/categorical/impute.py b/sdm/processing/categorical/impute.py index 166ec159b2..9e3b89d663 100644 --- a/sdm/processing/categorical/impute.py +++ b/sdm/processing/categorical/impute.py @@ -74,8 +74,8 @@ def _fit( fill_values.append(counts.argmax(dim=-1, keepdim=True)) self._set_buffer( - "_fill_values", - torch.stack(fill_values, dim=-1) + 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), ) diff --git a/sdm/processing/execution.py b/sdm/processing/execution.py index 317d04c46a..e0fa990762 100644 --- a/sdm/processing/execution.py +++ b/sdm/processing/execution.py @@ -75,8 +75,12 @@ def fit_transform( ) for module in series_processors: module._set_buffer( - "_offset", - torch.zeros(len(y), dtype=torch.long, device=y[0].device), + 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,9 +142,9 @@ def fit_transform( for module in processor.modules(): if isinstance(module, sp.SeriesEnsembleProcessor): module._set_buffer( - "_offset", - torch.tensor( - [ + name="_offset", + value=torch.tensor( + data=[ ensemble_table[i].size(-2) for i in range(len(ensemble_table)) ], @@ -184,9 +188,9 @@ def fit_transform( for module in series_processors: module._set_buffer( - "_offset", - torch.tensor( - [x[i].size(-2) for i in range(len(x))], + name="_offset", + value=torch.tensor( + data=[x[i].size(-2) for i in range(len(x))], dtype=torch.long, device=x[0].device, ), diff --git a/sdm/processing/numerical/impute.py b/sdm/processing/numerical/impute.py index 98a9099b77..c1d9105fc9 100644 --- a/sdm/processing/numerical/impute.py +++ b/sdm/processing/numerical/impute.py @@ -47,7 +47,8 @@ def _fit( keepdim=True, ) self._set_buffer( - "_mean", torch.where(mean.isnan(), self.fill_value, mean) + name="_mean", + value=torch.where(mean.isnan(), self.fill_value, mean), ) def _transform(self, table: TableTensor) -> TableTensor: diff --git a/sdm/processing/numerical/pca.py b/sdm/processing/numerical/pca.py index e1df411d7c..dfa7c91e28 100644 --- a/sdm/processing/numerical/pca.py +++ b/sdm/processing/numerical/pca.py @@ -53,7 +53,8 @@ def _fit( table.numerical.size(-1), ) self._set_buffer( - "components", vh[..., :num_components, :].transpose(-2, -1) + name="components", + value=vh[..., :num_components, :].transpose(-2, -1), ) def _transform(self, table: TableTensor) -> TableTensor: diff --git a/sdm/processing/numerical/power.py b/sdm/processing/numerical/power.py index 9c5e970e5f..22e5b40c39 100644 --- a/sdm/processing/numerical/power.py +++ b/sdm/processing/numerical/power.py @@ -278,8 +278,8 @@ def _fit( # ``mean`` never exceeds the column maximum, and is zero for an # entirely missing column, whose max must be zero. self._set_buffer( - "max", - torch.where(finite, table.numerical, mean).amax( + name="max", + value=torch.where(finite, table.numerical, mean).amax( dim=-2, keepdim=True, ), @@ -287,10 +287,10 @@ def _fit( del finite self._set_buffer( - "lambdas", - self._optimize_lambdas( - finite_or_nan, - constant_features, + name="lambdas", + value=self._optimize_lambdas( + inp=finite_or_nan, + constant_features=constant_features, count=count, ), ) From 018444f055ece9a6eb22b645631254318b2cbcdd Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Fri, 16 Oct 2026 01:30:30 +0900 Subject: [PATCH 4/9] Clarify why Standardize keeps its direct column check --- sdm/processing/numerical/standardize.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/sdm/processing/numerical/standardize.py b/sdm/processing/numerical/standardize.py index 1fc5e3e28a..d6e2ae60ab 100644 --- a/sdm/processing/numerical/standardize.py +++ b/sdm/processing/numerical/standardize.py @@ -52,7 +52,8 @@ def with_state(self, mean: Tensor | float, scale: Tensor | float) -> Self: return self def _handles_table(self, table: TableTensor) -> bool: - # PyTorch 2.7 cannot trace iteration over handled column types. + # PyTorch 2.7 fails fullgraph=True on the base class's frozenset loop. + # Standardize only needs to check whether numerical columns exist. return table.numerical.size(-1) > 0 def _fit( From e7cd281a62fdeeaf850783e32f3e705bdb50e639 Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Fri, 16 Oct 2026 01:31:09 +0900 Subject: [PATCH 5/9] Simplify Standardize compatibility comment --- sdm/processing/numerical/standardize.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/sdm/processing/numerical/standardize.py b/sdm/processing/numerical/standardize.py index d6e2ae60ab..dae941b9ef 100644 --- a/sdm/processing/numerical/standardize.py +++ b/sdm/processing/numerical/standardize.py @@ -52,8 +52,7 @@ def with_state(self, mean: Tensor | float, scale: Tensor | float) -> Self: return self def _handles_table(self, table: TableTensor) -> bool: - # PyTorch 2.7 fails fullgraph=True on the base class's frozenset loop. - # Standardize only needs to check whether numerical columns exist. + # The shared check fails with fullgraph=True on PyTorch 2.7. return table.numerical.size(-1) > 0 def _fit( From 672c1464ed72b4550c07273a1e74ede71b956647 Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Fri, 16 Oct 2026 01:34:39 +0900 Subject: [PATCH 6/9] Restore original Standardize compilation comment --- sdm/processing/numerical/standardize.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdm/processing/numerical/standardize.py b/sdm/processing/numerical/standardize.py index dae941b9ef..1fc5e3e28a 100644 --- a/sdm/processing/numerical/standardize.py +++ b/sdm/processing/numerical/standardize.py @@ -52,7 +52,7 @@ def with_state(self, mean: Tensor | float, scale: Tensor | float) -> Self: return self def _handles_table(self, table: TableTensor) -> bool: - # The shared check fails with fullgraph=True on PyTorch 2.7. + # PyTorch 2.7 cannot trace iteration over handled column types. return table.numerical.size(-1) > 0 def _fit( From 0597d44094b3692f19920b82da1345425b96b7c4 Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Fri, 16 Oct 2026 02:26:49 +0900 Subject: [PATCH 7/9] Use immutable tuples for handled column types --- sdm/processing/base.py | 2 +- sdm/processing/categorical/align.py | 2 +- sdm/processing/categorical/count.py | 2 +- sdm/processing/categorical/impute.py | 2 +- sdm/processing/categorical/shuffle.py | 2 +- sdm/processing/common/callable.py | 2 +- sdm/processing/common/choice.py | 8 ++++++-- sdm/processing/common/drop.py | 4 +++- sdm/processing/common/ensemble.py | 2 +- sdm/processing/common/identity.py | 2 +- sdm/processing/common/select.py | 2 +- sdm/processing/common/sequential.py | 10 ++++++---- sdm/processing/common/shuffle.py | 2 +- sdm/processing/common/stype.py | 4 +--- sdm/processing/common/table.py | 12 +++++++----- sdm/processing/common/task.py | 14 ++++++++------ sdm/processing/common/to_numerical.py | 2 +- sdm/processing/datetime/calendar.py | 2 +- sdm/processing/numerical/cast.py | 2 +- sdm/processing/numerical/clip.py | 2 +- sdm/processing/numerical/clip_soft.py | 2 +- sdm/processing/numerical/constant.py | 2 +- sdm/processing/numerical/detrend.py | 2 +- sdm/processing/numerical/flip_sign.py | 2 +- sdm/processing/numerical/impute.py | 2 +- sdm/processing/numerical/interpolate.py | 2 +- sdm/processing/numerical/pca.py | 2 +- sdm/processing/numerical/power.py | 2 +- sdm/processing/numerical/quantile.py | 2 +- sdm/processing/numerical/quantile_clip.py | 2 +- sdm/processing/numerical/random_projection.py | 2 +- sdm/processing/numerical/rank_gaussian.py | 2 +- sdm/processing/numerical/robust_scale.py | 2 +- sdm/processing/numerical/sigma_clip.py | 2 +- sdm/processing/numerical/standardize.py | 6 +----- sdm/processing/output/reduce.py | 2 +- sdm/processing/output/softmax.py | 2 +- sdm/processing/output/sort.py | 2 +- sdm/processing/text/sentence_transformer.py | 2 +- sdm/processing/text/tfidf.py | 2 +- test/models/test_base.py | 2 +- test/processing/common/test_choice.py | 4 ++-- test/processing/common/test_drop.py | 10 ++++++++-- test/processing/common/test_table.py | 2 +- test/processing/test_contract.py | 2 +- test/processing/test_ensemble.py | 4 ++-- test/processing/test_fitted_guard.py | 4 ++-- test/relational/test_sampler.py | 2 +- 48 files changed, 83 insertions(+), 71 deletions(-) diff --git a/sdm/processing/base.py b/sdm/processing/base.py index b6fd6a53ce..428d2a108d 100644 --- a/sdm/processing/base.py +++ b/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 diff --git a/sdm/processing/categorical/align.py b/sdm/processing/categorical/align.py index 6d146bca49..0d76bd808f 100644 --- a/sdm/processing/categorical/align.py +++ b/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__( diff --git a/sdm/processing/categorical/count.py b/sdm/processing/categorical/count.py index 1554291265..9dd41732fd 100644 --- a/sdm/processing/categorical/count.py +++ b/sdm/processing/categorical/count.py @@ -32,7 +32,7 @@ class AddCategoryCounts(Processor): vocabulary size. Defaults to 0. """ - handles_stypes = frozenset({Stype.categorical}) + handles_stypes = (Stype.categorical,) requires_fit = True _log_counts: Tensor diff --git a/sdm/processing/categorical/impute.py b/sdm/processing/categorical/impute.py index 9e3b89d663..c7d3bbc304 100644 --- a/sdm/processing/categorical/impute.py +++ b/sdm/processing/categorical/impute.py @@ -25,7 +25,7 @@ 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 diff --git a/sdm/processing/categorical/shuffle.py b/sdm/processing/categorical/shuffle.py index 8497d347fc..976e3f0f7f 100644 --- a/sdm/processing/categorical/shuffle.py +++ b/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__( diff --git a/sdm/processing/common/callable.py b/sdm/processing/common/callable.py index 679a4eee03..4427b78028 100644 --- a/sdm/processing/common/callable.py +++ b/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: diff --git a/sdm/processing/common/choice.py b/sdm/processing/common/choice.py index 25da3bac82..465d08826b 100644 --- a/sdm/processing/common/choice.py +++ b/sdm/processing/common/choice.py @@ -36,8 +36,12 @@ 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 = tuple( + dict.fromkeys( + stype + for option in self.options + for stype in option.handles_stypes + ) ) self.method = method self._option_ids: tuple[int, ...] = () diff --git a/sdm/processing/common/drop.py b/sdm/processing/common/drop.py index 340a0fd039..49be559476 100644 --- a/sdm/processing/common/drop.py +++ b/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) diff --git a/sdm/processing/common/ensemble.py b/sdm/processing/common/ensemble.py index 4f5a2cb5d3..e06e215413 100644 --- a/sdm/processing/common/ensemble.py +++ b/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 diff --git a/sdm/processing/common/identity.py b/sdm/processing/common/identity.py index 079e61bd82..23136a9ac2 100644 --- a/sdm/processing/common/identity.py +++ b/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: diff --git a/sdm/processing/common/select.py b/sdm/processing/common/select.py index aca6e9e73e..027948b6d1 100644 --- a/sdm/processing/common/select.py +++ b/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__( diff --git a/sdm/processing/common/sequential.py b/sdm/processing/common/sequential.py index b19fa74f32..56d5ad4def 100644 --- a/sdm/processing/common/sequential.py +++ b/sdm/processing/common/sequential.py @@ -26,12 +26,14 @@ 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 tuple( + dict.fromkeys( + stype for child in self for stype in child.handles_stypes + ) ) def append(self, processor: object) -> Self: diff --git a/sdm/processing/common/shuffle.py b/sdm/processing/common/shuffle.py index d349e8253b..26aace5227 100644 --- a/sdm/processing/common/shuffle.py +++ b/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__( diff --git a/sdm/processing/common/stype.py b/sdm/processing/common/stype.py index 32a57bc6fa..251f7faae5 100644 --- a/sdm/processing/common/stype.py +++ b/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() ) diff --git a/sdm/processing/common/table.py b/sdm/processing/common/table.py index c797a4b918..e8a26a2780 100644 --- a/sdm/processing/common/table.py +++ b/sdm/processing/common/table.py @@ -39,12 +39,14 @@ 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 tuple( + dict.fromkeys( + stype + for processor in self.processors.values() + for stype in processor.handles_stypes + ) ) def _fit_ensemble( diff --git a/sdm/processing/common/task.py b/sdm/processing/common/task.py index f163976855..4cccf32cbd 100644 --- a/sdm/processing/common/task.py +++ b/sdm/processing/common/task.py @@ -43,16 +43,18 @@ 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 tuple( + dict.fromkeys( + stype + for processor in self.processors.values() + for stype in processor.handles_stypes + ) ) if self._task not in self.processors: - return frozenset() + return () return self.processors[self._task].handles_stypes def _processor(self) -> EnsembleProcessor: diff --git a/sdm/processing/common/to_numerical.py b/sdm/processing/common/to_numerical.py index 064aeb74fb..ce745914fb 100644 --- a/sdm/processing/common/to_numerical.py +++ b/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``.""" diff --git a/sdm/processing/datetime/calendar.py b/sdm/processing/datetime/calendar.py index 4a4a4e97f9..7590788e46 100644 --- a/sdm/processing/datetime/calendar.py +++ b/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__( diff --git a/sdm/processing/numerical/cast.py b/sdm/processing/numerical/cast.py index 7ae8014648..4a32cdc81f 100644 --- a/sdm/processing/numerical/cast.py +++ b/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: diff --git a/sdm/processing/numerical/clip.py b/sdm/processing/numerical/clip.py index fbb0ce753f..1ca630ed9f 100644 --- a/sdm/processing/numerical/clip.py +++ b/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__( diff --git a/sdm/processing/numerical/clip_soft.py b/sdm/processing/numerical/clip_soft.py index efda764ee0..8d403d1cb2 100644 --- a/sdm/processing/numerical/clip_soft.py +++ b/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: diff --git a/sdm/processing/numerical/constant.py b/sdm/processing/numerical/constant.py index 4fe2873cf3..adc811b85e 100644 --- a/sdm/processing/numerical/constant.py +++ b/sdm/processing/numerical/constant.py @@ -27,7 +27,7 @@ class DropConstantColumns(EnsembleProcessor): Must be positive. """ - handles_stypes = frozenset({Stype.numerical}) + handles_stypes = (Stype.numerical,) requires_fit = True def __init__( diff --git a/sdm/processing/numerical/detrend.py b/sdm/processing/numerical/detrend.py index 92a7c0b0ca..0b81b7aca9 100644 --- a/sdm/processing/numerical/detrend.py +++ b/sdm/processing/numerical/detrend.py @@ -19,7 +19,7 @@ class Detrend(SeriesEnsembleProcessor, EnsembleInvertibleMixin): deviation before detrending. """ - handles_stypes = frozenset({Stype.numerical}) + handles_stypes = (Stype.numerical,) requires_fit = True def __init__(self, threshold: float | None = None) -> None: diff --git a/sdm/processing/numerical/flip_sign.py b/sdm/processing/numerical/flip_sign.py index d7ce154b08..57cec8958c 100644 --- a/sdm/processing/numerical/flip_sign.py +++ b/sdm/processing/numerical/flip_sign.py @@ -16,7 +16,7 @@ class FlipSign(EnsembleProcessor, EnsembleInvertibleMixin): probability: Probability of negating each numerical column. """ - handles_stypes = frozenset({Stype.numerical}) + handles_stypes = (Stype.numerical,) requires_fit = True def __init__( diff --git a/sdm/processing/numerical/impute.py b/sdm/processing/numerical/impute.py index c1d9105fc9..01e486692d 100644 --- a/sdm/processing/numerical/impute.py +++ b/sdm/processing/numerical/impute.py @@ -20,7 +20,7 @@ class ImputeMean(Processor): undefined (columns without finite values). """ - handles_stypes = frozenset({Stype.numerical}) + handles_stypes = (Stype.numerical,) requires_fit = True _mean: Tensor diff --git a/sdm/processing/numerical/interpolate.py b/sdm/processing/numerical/interpolate.py index 5fa61fba98..aa13fcce44 100644 --- a/sdm/processing/numerical/interpolate.py +++ b/sdm/processing/numerical/interpolate.py @@ -17,7 +17,7 @@ class Interpolate(SeriesEnsembleProcessor, EnsembleInvertibleMixin): fill_value: Value used for columns with no observations. """ - handles_stypes = frozenset({Stype.numerical}) + handles_stypes = (Stype.numerical,) requires_fit = True def __init__(self, *, fill_value: float = 0.0) -> None: diff --git a/sdm/processing/numerical/pca.py b/sdm/processing/numerical/pca.py index dfa7c91e28..2924ec561f 100644 --- a/sdm/processing/numerical/pca.py +++ b/sdm/processing/numerical/pca.py @@ -17,7 +17,7 @@ class PCA(Processor): num_components: Number of principal components to keep. """ - handles_stypes = frozenset({Stype.numerical}) + handles_stypes = (Stype.numerical,) requires_fit = True mean: Tensor diff --git a/sdm/processing/numerical/power.py b/sdm/processing/numerical/power.py index 22e5b40c39..aa7709daee 100644 --- a/sdm/processing/numerical/power.py +++ b/sdm/processing/numerical/power.py @@ -226,7 +226,7 @@ class PowerTransform(Processor, InvertibleMixin): features using statistics fitted after the power transform. """ - handles_stypes = frozenset({Stype.numerical}) + handles_stypes = (Stype.numerical,) requires_fit = True lambdas: Tensor diff --git a/sdm/processing/numerical/quantile.py b/sdm/processing/numerical/quantile.py index e3f3f19bad..e9c276fc1b 100644 --- a/sdm/processing/numerical/quantile.py +++ b/sdm/processing/numerical/quantile.py @@ -93,7 +93,7 @@ class QuantileTransform(Processor, InvertibleMixin): output_distribution: Distribution to map the empirical quantiles to. """ - handles_stypes = frozenset({Stype.numerical}) + handles_stypes = (Stype.numerical,) requires_fit = True _quantiles: Tensor diff --git a/sdm/processing/numerical/quantile_clip.py b/sdm/processing/numerical/quantile_clip.py index e8305d6ea5..f8f64775a4 100644 --- a/sdm/processing/numerical/quantile_clip.py +++ b/sdm/processing/numerical/quantile_clip.py @@ -21,7 +21,7 @@ class ClipQuantiles(Processor): bound. Must satisfy ``0 <= q_low <= q_high <= 1``. """ - handles_stypes = frozenset({Stype.numerical}) + handles_stypes = (Stype.numerical,) requires_fit = True lower_bound: Tensor diff --git a/sdm/processing/numerical/random_projection.py b/sdm/processing/numerical/random_projection.py index 7558d7151d..2ad600422a 100644 --- a/sdm/processing/numerical/random_projection.py +++ b/sdm/processing/numerical/random_projection.py @@ -20,7 +20,7 @@ class RandomProjection(EnsembleProcessor): ``1 / sqrt(channels)``. """ - handles_stypes = frozenset({Stype.numerical}) + handles_stypes = (Stype.numerical,) requires_fit = True def __init__( diff --git a/sdm/processing/numerical/rank_gaussian.py b/sdm/processing/numerical/rank_gaussian.py index 62eef9b72d..fd95a19eb0 100644 --- a/sdm/processing/numerical/rank_gaussian.py +++ b/sdm/processing/numerical/rank_gaussian.py @@ -32,7 +32,7 @@ class RankGaussian(Processor): ``None``, retain all fitted values. """ - handles_stypes = frozenset({Stype.numerical}) + handles_stypes = (Stype.numerical,) requires_fit = True _values: Tensor diff --git a/sdm/processing/numerical/robust_scale.py b/sdm/processing/numerical/robust_scale.py index 2043e37d95..5edc2ff83a 100644 --- a/sdm/processing/numerical/robust_scale.py +++ b/sdm/processing/numerical/robust_scale.py @@ -22,7 +22,7 @@ class RobustScale(Processor, InvertibleMixin): used as the scale. """ - handles_stypes = frozenset({Stype.numerical}) + handles_stypes = (Stype.numerical,) requires_fit = True median: Tensor diff --git a/sdm/processing/numerical/sigma_clip.py b/sdm/processing/numerical/sigma_clip.py index f477d685dd..f8c3a2594c 100644 --- a/sdm/processing/numerical/sigma_clip.py +++ b/sdm/processing/numerical/sigma_clip.py @@ -26,7 +26,7 @@ class ClipSigma(Processor): deviations from the mean mark the soft clipping bounds. """ - handles_stypes = frozenset({Stype.numerical}) + handles_stypes = (Stype.numerical,) requires_fit = True lower_bound: Tensor diff --git a/sdm/processing/numerical/standardize.py b/sdm/processing/numerical/standardize.py index 1fc5e3e28a..92827b64e3 100644 --- a/sdm/processing/numerical/standardize.py +++ b/sdm/processing/numerical/standardize.py @@ -25,7 +25,7 @@ class Standardize(Processor, InvertibleMixin): eps: Value added to each fitted standard deviation. """ - handles_stypes = frozenset({Stype.numerical}) + handles_stypes = (Stype.numerical,) requires_fit = True mean: Tensor @@ -51,10 +51,6 @@ def with_state(self, mean: Tensor | float, scale: Tensor | float) -> Self: self.requires_fit = False return self - def _handles_table(self, table: TableTensor) -> bool: - # PyTorch 2.7 cannot trace iteration over handled column types. - return table.numerical.size(-1) > 0 - def _fit( self, table: TableTensor, diff --git a/sdm/processing/output/reduce.py b/sdm/processing/output/reduce.py index 4ea8d4ae18..b187da5d62 100644 --- a/sdm/processing/output/reduce.py +++ b/sdm/processing/output/reduce.py @@ -22,7 +22,7 @@ class AverageEstimators(EnsembleProcessor): aggregation less sensitive to outlying estimator predictions. """ - handles_stypes = frozenset({Stype.numerical}) + handles_stypes = (Stype.numerical,) requires_fit = False def __init__( diff --git a/sdm/processing/output/softmax.py b/sdm/processing/output/softmax.py index c29ef71756..3305e4b182 100644 --- a/sdm/processing/output/softmax.py +++ b/sdm/processing/output/softmax.py @@ -18,7 +18,7 @@ class Softmax(Processor): higher values produce a softer distribution. """ - handles_stypes = frozenset({Stype.numerical}) + handles_stypes = (Stype.numerical,) requires_fit = False def __init__( diff --git a/sdm/processing/output/sort.py b/sdm/processing/output/sort.py index 363b118c16..9990936f5d 100644 --- a/sdm/processing/output/sort.py +++ b/sdm/processing/output/sort.py @@ -14,7 +14,7 @@ class SortQuantiles(Processor): If ``None``, all numerical columns form one quantile group. """ - handles_stypes = frozenset({Stype.numerical}) + handles_stypes = (Stype.numerical,) requires_fit = False def __init__(self, num_quantiles: int | None = None) -> None: diff --git a/sdm/processing/text/sentence_transformer.py b/sdm/processing/text/sentence_transformer.py index eaf6687890..deee048621 100644 --- a/sdm/processing/text/sentence_transformer.py +++ b/sdm/processing/text/sentence_transformer.py @@ -400,7 +400,7 @@ class SentenceTransformer(Processor): """ requires_fit = False - handles_stypes = frozenset({Stype.text}) + handles_stypes = (Stype.text,) def __init__( self, diff --git a/sdm/processing/text/tfidf.py b/sdm/processing/text/tfidf.py index cde22ea392..08a920b272 100644 --- a/sdm/processing/text/tfidf.py +++ b/sdm/processing/text/tfidf.py @@ -46,7 +46,7 @@ class TFIDF(EnsembleProcessor): lowercase: If ``True``, lowercase text before tokenizing. """ - handles_stypes = frozenset({Stype.text}) + handles_stypes = (Stype.text,) requires_fit = True _idf_template: Tensor diff --git a/test/models/test_base.py b/test/models/test_base.py index 91f1784969..581c43cb73 100644 --- a/test/models/test_base.py +++ b/test/models/test_base.py @@ -181,7 +181,7 @@ def on_model_forward_end( class _GeneratorRecordingProcessor(Processor, InvertibleMixin): - handles_stypes = frozenset(Stype) + handles_stypes = tuple(Stype) requires_fit = True generators: ClassVar[list[torch.Generator | None]] = [] draws: ClassVar[list[torch.Tensor]] = [] diff --git a/test/processing/common/test_choice.py b/test/processing/common/test_choice.py index e2b75c5b73..aeec54426f 100644 --- a/test/processing/common/test_choice.py +++ b/test/processing/common/test_choice.py @@ -17,7 +17,7 @@ class Add(Processor, InvertibleMixin): requires_fit = False - handles_stypes = frozenset(Stype) + handles_stypes = tuple(Stype) def __init__(self, value: float) -> None: super().__init__() @@ -31,7 +31,7 @@ def _inverse_transform(self, table: TableTensor) -> TableTensor: class AddFittedMemberCount(EnsembleProcessor): - handles_stypes = frozenset(Stype) + handles_stypes = tuple(Stype) requires_fit = True def __init__(self) -> None: diff --git a/test/processing/common/test_drop.py b/test/processing/common/test_drop.py index 0535fd7afd..d385e31040 100644 --- a/test/processing/common/test_drop.py +++ b/test/processing/common/test_drop.py @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import pytest import torch import sdm.processing as sp @@ -17,10 +18,15 @@ def _mixed_table() -> TableTensor: ) -def test_drop_stypes_removes_configured_stypes() -> None: +@pytest.mark.parametrize( + "stypes", [(Stype.numerical,), ("numerical", Stype.numerical)] +) +def test_drop_stypes_removes_configured_stypes( + stypes: tuple[Stype | str, ...], +) -> None: table = _mixed_table() - output = sp.DropStypes(Stype.numerical).fit_transform(table) + output = sp.DropStypes(*stypes).fit_transform(table) assert output.columns[Stype.numerical] == () assert output.columns[Stype.categorical] == ("cat_0",) diff --git a/test/processing/common/test_table.py b/test/processing/common/test_table.py index a747374a47..1a32fd8c93 100644 --- a/test/processing/common/test_table.py +++ b/test/processing/common/test_table.py @@ -11,7 +11,7 @@ class _Add(Processor): requires_fit = False - handles_stypes = frozenset({Stype.numerical}) + handles_stypes = (Stype.numerical,) def __init__(self, value: float) -> None: super().__init__() diff --git a/test/processing/test_contract.py b/test/processing/test_contract.py index 90ba5f4db0..ca571219a0 100644 --- a/test/processing/test_contract.py +++ b/test/processing/test_contract.py @@ -286,7 +286,7 @@ def test_preserves_rows_and_unhandled_stypes( before = table[member_id] after = output[member_id] assert after.size()[:-1] == before.size()[:-1] - for stype in before.active_stypes - processor.handles_stypes: + for stype in before.active_stypes.difference(processor.handles_stypes): columns = before.columns[stype] assert all(column in after.column_names for column in columns) assert all(after.stype(column) == stype for column in columns) diff --git a/test/processing/test_ensemble.py b/test/processing/test_ensemble.py index 177f480e88..b8d555bfc5 100644 --- a/test/processing/test_ensemble.py +++ b/test/processing/test_ensemble.py @@ -20,7 +20,7 @@ # Generator forwarding is only observable through a stochastic transformation # and is therefore left to those processors as well. class IdentityEnsembleProcessor(EnsembleProcessor): - handles_stypes = frozenset({Stype.numerical}) + handles_stypes = (Stype.numerical,) requires_fit = True def _fit_ensemble( @@ -65,7 +65,7 @@ def _inverse_transform_ensemble( class _StatelessProcessor(Processor, InvertibleMixin): - handles_stypes = frozenset({Stype.numerical}) + handles_stypes = (Stype.numerical,) requires_fit = False def _transform(self, table: TableTensor) -> TableTensor: diff --git a/test/processing/test_fitted_guard.py b/test/processing/test_fitted_guard.py index 762d233f0e..164d4d3a64 100644 --- a/test/processing/test_fitted_guard.py +++ b/test/processing/test_fitted_guard.py @@ -54,7 +54,7 @@ def test_invertible_processor_requires_fit_for_inverse_transform( class StatelessProcessor(Processor): - handles_stypes = frozenset({Stype.numerical}) + handles_stypes = (Stype.numerical,) requires_fit = False def _transform(self, table: TableTensor) -> TableTensor: @@ -62,7 +62,7 @@ def _transform(self, table: TableTensor) -> TableTensor: class StatefulProcessor(Processor): - handles_stypes = frozenset({Stype.numerical}) + handles_stypes = (Stype.numerical,) requires_fit = True def _fit( diff --git a/test/relational/test_sampler.py b/test/relational/test_sampler.py index 39b9258fd7..9d954487a3 100644 --- a/test/relational/test_sampler.py +++ b/test/relational/test_sampler.py @@ -209,7 +209,7 @@ def test_batched_recipe_execution(relational_data: RelationalData) -> None: pytest.importorskip("pyg_lib") class ShapeLock(Processor): - handles_stypes = frozenset({Stype.numerical}) + handles_stypes = (Stype.numerical,) requires_fit = True def _fit( From 0718a57ebdfe5d253e6f65dbf914ec37ca51b861 Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Fri, 16 Oct 2026 02:40:48 +0900 Subject: [PATCH 8/9] Explain tracing rationale for processor column checks --- sdm/processing/base.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/sdm/processing/base.py b/sdm/processing/base.py index 428d2a108d..3f34a51fe4 100644 --- a/sdm/processing/base.py +++ b/sdm/processing/base.py @@ -161,7 +161,8 @@ def _fit_transform( return self._transform(table) def _handles_table(self, table: TableTensor) -> bool: - # Check handled blocks directly instead of constructing active types. + # 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 From 4f4ca60fd6414fe42d948ac65569a4349c986b35 Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Fri, 16 Oct 2026 03:28:12 +0900 Subject: [PATCH 9/9] Share ordered column-type merging across composite processors --- sdm/processing/common/choice.py | 8 +------- sdm/processing/common/sequential.py | 6 +----- sdm/processing/common/table.py | 8 +------- sdm/processing/common/task.py | 8 +------- sdm/processing/ensemble.py | 14 +++++++++++++- 5 files changed, 17 insertions(+), 27 deletions(-) diff --git a/sdm/processing/common/choice.py b/sdm/processing/common/choice.py index 465d08826b..a384452690 100644 --- a/sdm/processing/common/choice.py +++ b/sdm/processing/common/choice.py @@ -36,13 +36,7 @@ def __init__( for arg in args: options.append(EnsembleProcessor.as_processor(arg)) self.options: ModuleList[EnsembleProcessor] = ModuleList(options) - self.handles_stypes = tuple( - dict.fromkeys( - 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, ...] = () diff --git a/sdm/processing/common/sequential.py b/sdm/processing/common/sequential.py index 56d5ad4def..e2576f2a7c 100644 --- a/sdm/processing/common/sequential.py +++ b/sdm/processing/common/sequential.py @@ -30,11 +30,7 @@ def handles_stypes(self) -> tuple[Stype, ...]: r""":meta private:""" # noqa: D415 if len(self) == 0: return tuple(Stype) - return tuple( - dict.fromkeys( - stype for child in self for stype in child.handles_stypes - ) - ) + return self._merge_stypes(self) def append(self, processor: object) -> Self: r"""Append a processor or callable to this sequence. diff --git a/sdm/processing/common/table.py b/sdm/processing/common/table.py index e8a26a2780..ced21cadfe 100644 --- a/sdm/processing/common/table.py +++ b/sdm/processing/common/table.py @@ -41,13 +41,7 @@ def __init__( @property def handles_stypes(self) -> tuple[Stype, ...]: r""":meta private:""" # noqa: D415 - return tuple( - dict.fromkeys( - stype - for processor in self.processors.values() - for stype in processor.handles_stypes - ) - ) + return self._merge_stypes(self.processors.values()) def _fit_ensemble( self, diff --git a/sdm/processing/common/task.py b/sdm/processing/common/task.py index 4cccf32cbd..a6bbc439c8 100644 --- a/sdm/processing/common/task.py +++ b/sdm/processing/common/task.py @@ -46,13 +46,7 @@ def __init__( def handles_stypes(self) -> tuple[Stype, ...]: r""":meta private:""" # noqa: D415 if self._task is None: - return tuple( - dict.fromkeys( - 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 () return self.processors[self._task].handles_stypes diff --git a/sdm/processing/ensemble.py b/sdm/processing/ensemble.py index 296916f9cc..9ece89286e 100644 --- a/sdm/processing/ensemble.py +++ b/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,