From d816644a539557a20aa56b245827459b6236a20b Mon Sep 17 00:00:00 2001 From: Valter Hudovernik Date: Sat, 26 Sep 2026 04:03:48 -0700 Subject: [PATCH 1/3] Add rank-Gaussian views and balanced class shifts to Kumo recipes --- sdm/models/kumo/tabular/recipe.py | 29 +++- sdm/processing/__init__.py | 2 + sdm/processing/categorical/shuffle.py | 40 ++++-- sdm/processing/numerical/__init__.py | 2 + sdm/processing/numerical/rank_gaussian.py | 72 ++++++++++ test/models/kumo/tabular/test_recipe.py | 70 ++++++++++ test/processing/categorical/test_shuffle.py | 50 ++++++- .../numerical/test_rank_gaussian.py | 132 ++++++++++++++++++ test/processing/test_contract.py | 2 + 9 files changed, 384 insertions(+), 15 deletions(-) create mode 100644 sdm/processing/numerical/rank_gaussian.py create mode 100644 test/processing/numerical/test_rank_gaussian.py diff --git a/sdm/models/kumo/tabular/recipe.py b/sdm/models/kumo/tabular/recipe.py index 38173f3b4..79c723099 100644 --- a/sdm/models/kumo/tabular/recipe.py +++ b/sdm/models/kumo/tabular/recipe.py @@ -31,7 +31,32 @@ def numerical_processor() -> sp.Sequential: features=[ sp.StypeDispatch( numerical=[ - numerical_processor(), + sp.Cast(torch.float64), + sp.DropConstantColumns(), + # Period 12 preserves the original three-way schedule + # while replacing every fourth view with Gaussian ranks. + sp.Choice( + *[ + [ + sp.RankGaussian(), + sp.Standardize(), + sp.ClipSigma(threshold=4.0), + ] + if i % 4 == 3 + else [ + sp.Standardize(eps=1e-6), + sp.Clip(-100.0, 100.0), + ( + sp.Identity(), + sp.PowerTransform(), + [sp.RobustScale(), sp.ClipSoft(3.0)], + )[i % 3], + sp.ClipSigma(threshold=4.0), + ] + for i in range(12) + ], + method="round_robin", + ), sp.FlipSign(), ], categorical=[ @@ -49,7 +74,7 @@ def numerical_processor() -> sp.Sequential: sp.StypeDispatch( categorical=[ sp.AlignCategories(), - sp.ShuffleCategories(method="shift"), + sp.ShuffleCategories(method="balanced_shift"), ], numerical=[ sp.Standardize(), diff --git a/sdm/processing/__init__.py b/sdm/processing/__init__.py index 400a11d9a..ad88657d6 100644 --- a/sdm/processing/__init__.py +++ b/sdm/processing/__init__.py @@ -32,6 +32,7 @@ ImputeMean, PowerTransform, QuantileTransform, + RankGaussian, Standardize, RobustScale, FlipSign, @@ -76,6 +77,7 @@ "ImputeMean", "PowerTransform", "QuantileTransform", + "RankGaussian", "Standardize", "RobustScale", "FlipSign", diff --git a/sdm/processing/categorical/shuffle.py b/sdm/processing/categorical/shuffle.py index 7b160c46e..f9d6b0123 100644 --- a/sdm/processing/categorical/shuffle.py +++ b/sdm/processing/categorical/shuffle.py @@ -12,7 +12,7 @@ class ShuffleCategories(EnsembleProcessor): - """Independently permute the integer codes of categorical columns. + """Permute the integer codes of categorical columns. One permutation per categorical column is drawn when the processor is fitted. Codes and their corresponding category vectors are permuted @@ -24,7 +24,10 @@ class ShuffleCategories(EnsembleProcessor): Args: method: Permutation strategy. ``"shift"`` cyclically shifts the codes by a drawn offset, and ``"random"`` remaps the codes with - a drawn permutation. + a drawn permutation. ``"balanced_shift"`` draws cyclic offsets + without replacement before starting another cycle, balancing + their counts across ensemble members for each column and class + count. """ handles_stypes = frozenset({Stype.categorical}) @@ -32,7 +35,7 @@ class ShuffleCategories(EnsembleProcessor): def __init__( self, - method: Literal["shift", "random"] = "random", + method: Literal["shift", "random", "balanced_shift"] = "random", ) -> None: super().__init__() self.method = method @@ -51,21 +54,34 @@ def _draw_permutations( self, table: TableTensor, *, + shifts: dict[tuple[torch.device, int, int], Tensor], generator: torch.Generator | None = None, ) -> list[Tensor]: device = table.categorical.device permutations: list[Tensor] = [] - for category in table.categorical.categories: + for column, category in enumerate(table.categorical.categories): n_classes = category.numel() if n_classes <= 1: permutation = torch.arange(n_classes, device=device) - elif self.method == "shift": - offset = torch.randint( - n_classes, - (1,), - generator=generator, - device=device, - ) + elif self.method in ("shift", "balanced_shift"): + if self.method == "balanced_shift": + key = (device, column, n_classes) + remaining = shifts.get(key) + if remaining is None or remaining.numel() == 0: + remaining = torch.randperm( + n_classes, + generator=generator, + device=device, + ) + offset = remaining[:1] + shifts[key] = remaining[1:] + else: + offset = torch.randint( + n_classes, + (1,), + generator=generator, + device=device, + ) permutation = ( torch.arange(n_classes, device=device) - offset ) % n_classes @@ -92,9 +108,11 @@ def _fit_ensemble( tuple[torch.device, tuple[tuple[int, ...], ...]], int ] = {} + shifts: dict[tuple[torch.device, int, int], Tensor] = {} for member_id in range(len(ensemble_table)): permutations = self._draw_permutations( ensemble_table[member_id], + shifts=shifts, generator=generator, ) key = ( diff --git a/sdm/processing/numerical/__init__.py b/sdm/processing/numerical/__init__.py index 6f670deb4..fbbba0ec8 100644 --- a/sdm/processing/numerical/__init__.py +++ b/sdm/processing/numerical/__init__.py @@ -11,6 +11,7 @@ from sdm.processing.numerical.impute import ImputeMean from sdm.processing.numerical.power import PowerTransform from sdm.processing.numerical.quantile import QuantileTransform +from sdm.processing.numerical.rank_gaussian import RankGaussian from sdm.processing.numerical.standardize import Standardize from sdm.processing.numerical.robust_scale import RobustScale from sdm.processing.numerical.flip_sign import FlipSign @@ -27,6 +28,7 @@ "ImputeMean", "PowerTransform", "QuantileTransform", + "RankGaussian", "Standardize", "RobustScale", "FlipSign", diff --git a/sdm/processing/numerical/rank_gaussian.py b/sdm/processing/numerical/rank_gaussian.py new file mode 100644 index 000000000..34bbb4ce4 --- /dev/null +++ b/sdm/processing/numerical/rank_gaussian.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import torch + +from sdm import Stype, TableTensor +from sdm.processing import Processor +from sdm.processing.numerical.quantile import _batched_interp + + +class RankGaussian(Processor): + """Map interpolated empirical mid-ranks to standard normal quantiles. + + Each fitted value receives probability ``(L + R) / (2 * N)``, where + ``L`` and ``R`` count finite fitted values strictly below and at or below + it, and ``N`` is the number of finite fitted values. Ties share a rank. + Query probabilities interpolate between fitted values and clamp to the + endpoint probabilities, keeping finite outputs even outside the range. + + NaN and infinite values are ignored during fitting. NaNs are preserved + during transformation. Constant columns map to zero; columns without + finite fitted values produce NaNs. All fitted values are retained. + """ + + handles_stypes = frozenset({Stype.numerical}) + requires_fit = True + + def __init__(self) -> None: + super().__init__() + self.register_buffer("_values", torch.empty(0)) + self.register_buffer("_probabilities", torch.empty(0)) + + def _fit( + self, + table: TableTensor, + *, + generator: torch.Generator | None = None, + ) -> None: + # [..., N, C] -> [..., C, N]; double precision keeps tail ranks open. + numerical = table.numerical.double().movedim(-1, -2) + values = numerical.masked_fill(~numerical.isfinite(), torch.inf) + values = values.sort(dim=-1).values.contiguous() + finite = values.isfinite() + count = finite.sum(dim=-1, keepdim=True) + left = torch.searchsorted(values, values, right=False) + right = torch.searchsorted(values, values, right=True) + probabilities = (left + right).double() / (2 * count.clamp_min(1)) + + # Pad missing observations with the last finite knot and its rank. + last = (count - 1).clamp_min(0) + self._values = torch.where(finite, values, values.gather(-1, last)) + self._probabilities = torch.where( + finite, probabilities, probabilities.gather(-1, last) + ) + self._values.masked_fill_(count == 0, torch.nan) + self._probabilities.masked_fill_(count == 0, torch.nan) + + def _transform(self, table: TableTensor) -> TableTensor: + numerical = table.numerical + columns = numerical.movedim(-1, -2) + n_rows = columns.size(-1) + n_fitted = self._values.size(-1) + probabilities = _batched_interp( + columns.reshape(-1, n_rows).to(self._values.dtype).contiguous(), + self._values.reshape(-1, n_fitted), + self._probabilities.reshape(-1, n_fitted), + ) + output = torch.special.ndtri(probabilities).reshape(columns.shape) + output = output.masked_fill(columns.isnan(), torch.nan) + return table.replace_blocks( + numerical=output.movedim(-1, -2).to(numerical.dtype) + ) diff --git a/test/models/kumo/tabular/test_recipe.py b/test/models/kumo/tabular/test_recipe.py index a8a455719..4fdff2a07 100644 --- a/test/models/kumo/tabular/test_recipe.py +++ b/test/models/kumo/tabular/test_recipe.py @@ -150,3 +150,73 @@ def test_default_recipe_reduces_outputs_per_task() -> None: ) assert output.size() == (5, 3) torch.testing.assert_close(output.numerical.sum(dim=-1), torch.ones(5)) + + +@withCUDA +@pytest.mark.parametrize("num_members", [1, 4, 8, 16]) +def test_default_recipe_numerical_schedule( + num_members: int, device: torch.device +) -> None: + values = torch.arange(1, 101, device=device).double().square() + values[7] = torch.nan + features = TableTensor.from_tensor(values.unsqueeze(-1)) + recipe = KumoTabular.default_recipe() + output = recipe.features.fit_transform_ensemble( + EnsembleTable.from_table(features, num_members=num_members) + ) + query = features[:13] + prediction_features = recipe.features.transform_ensemble( + EnsembleTable.from_table(query, num_members=num_members) + ) + for i in range(num_members): + if i % 4 == 3: + processor = sp.Sequential( + sp.RankGaussian(), + sp.Standardize(), + sp.ClipSigma(threshold=4.0), + ) + else: + processor = sp.Sequential( + sp.Standardize(eps=1e-6), + sp.Clip(-100.0, 100.0), + ( + sp.Identity(), + sp.PowerTransform(), + [sp.RobustScale(), sp.ClipSoft(3.0)], + )[i % 3], + sp.ClipSigma(threshold=4.0), + ) + expected = processor.fit_transform(features).numerical.float() + # Independent sign flips are allowed; magnitudes identify each view. + torch.testing.assert_close( + output[i].numerical.abs(), expected.abs(), equal_nan=True + ) + torch.testing.assert_close( + prediction_features[i].numerical, + output[i].numerical[:13], + equal_nan=True, + ) + + +@withCUDA +@pytest.mark.parametrize("num_members", [1, 8, 16]) +@pytest.mark.parametrize("num_classes", [2, 3, 10]) +def test_default_recipe_balances_class_shifts( + num_members: int, num_classes: int, device: torch.device +) -> None: + target = TableTensor( + categorical=CategoricalTensor( + code=torch.arange(num_classes, device=device).unsqueeze(-1), + categories=(torch.arange(num_classes, device=device),), + ) + ) + output = KumoTabular.default_recipe().target.fit_transform_ensemble( + EnsembleTable.from_table(target, num_members=num_members) + ) + codes = torch.stack( + [output[i].categorical.code[0, 0] for i in range(num_members)] + ) + counts = codes.long().bincount(minlength=num_classes) + assert counts.max() - counts.min() <= 1 + for i in range(num_members): + assert output[i].categorical.tolist() == target.categorical.tolist() diff --git a/test/processing/categorical/test_shuffle.py b/test/processing/categorical/test_shuffle.py index 392e428bc..6757830f5 100644 --- a/test/processing/categorical/test_shuffle.py +++ b/test/processing/categorical/test_shuffle.py @@ -84,9 +84,9 @@ def test_shuffle_categories_random_permutes_each_categorical_column( assert transformed.categorical.tolist() == features.categorical.tolist() -@pytest.mark.parametrize("method", ["shift", "random"]) +@pytest.mark.parametrize("method", ["shift", "random", "balanced_shift"]) def test_shuffle_categories_is_reproducible_with_generator( - method: Literal["shift", "random"], + method: Literal["shift", "random", "balanced_shift"], ) -> None: features = _table( [[0, 0], [1, 1], [2, -1], [1, 0]], @@ -201,3 +201,49 @@ def test_shuffle_categories_refit_replaces_ensemble_state() -> None: ) assert output.equal(expected) + + +@withCUDA +@pytest.mark.parametrize("num_classes", [1, 2, 3, 10]) +@pytest.mark.parametrize("num_members", [1, 8, 21]) +def test_balanced_shifts_cover_classes( + device: torch.device, + num_classes: int, + num_members: int, +) -> None: + target = _table( + [[i] for i in range(num_classes)] + [[-1]], + (tuple(str(i) for i in range(num_classes)),), + device=device, + ) + ensemble = EnsembleTable.from_table(target, num_members=num_members) + processor = ShuffleCategories(method="balanced_shift") + output = processor.fit_transform_ensemble(ensemble) + shifts = torch.stack( + [output[i].categorical.code[0, 0] for i in range(num_members)] + ).long() + counts = shifts.bincount(minlength=num_classes) + assert int(counts.max() - counts.min()) <= 1 + for start in range(0, num_members, num_classes): + cycle = shifts[start : start + num_classes] + assert cycle.unique().numel() == cycle.numel() + query_output = processor.transform_ensemble(ensemble) + for i in range(num_members): + assert output[i].categorical.tolist() == target.categorical.tolist() + assert output[i].equal(query_output[i]) + + +def test_balanced_shifts_reproducible_ensemble() -> None: + target = _table([[0], [1], [2], [-1]], (("a", "b", "c"),)) + ensemble = EnsembleTable.from_table(target, num_members=8) + processor = ShuffleCategories(method="balanced_shift") + first = processor.fit_transform_ensemble( + ensemble, + generator=torch.Generator().manual_seed(7), + ) + second = processor.fit_transform_ensemble( + ensemble, + generator=torch.Generator().manual_seed(7), + ) + for i in range(8): + assert first[i].equal(second[i]) diff --git a/test/processing/numerical/test_rank_gaussian.py b/test/processing/numerical/test_rank_gaussian.py new file mode 100644 index 000000000..f71fafcb3 --- /dev/null +++ b/test/processing/numerical/test_rank_gaussian.py @@ -0,0 +1,132 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + +from sdm import TableTensor +from sdm.processing import RankGaussian +from sdm.testing import withCUDA + + +@withCUDA +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +def test_mid_ranks_and_interpolated_query( + dtype: torch.dtype, device: torch.device +) -> None: + context = torch.tensor( + [[0.0], [1.0], [1.0], [2.0]], dtype=dtype, device=device + ) + processor = RankGaussian().fit(TableTensor.from_tensor(context)) + probabilities = context.new_tensor([[0.125], [0.5], [0.5], [0.875]]) + torch.testing.assert_close( + processor.transform(TableTensor.from_tensor(context)).numerical, + torch.special.ndtri(probabilities), + ) + query = context.new_tensor([[-100.0], [0.5], [1.5], [100.0], [torch.nan]]) + expected = context.new_tensor( + [[0.125], [0.3125], [0.6875], [0.875], [torch.nan]] + ) + torch.testing.assert_close( + processor.transform(TableTensor.from_tensor(query)).numerical, + torch.special.ndtri(expected), + equal_nan=True, + ) + + +@withCUDA +def test_ties_at_endpoints_use_mid_ranks(device: torch.device) -> None: + context = torch.tensor([[0.0], [0.0], [2.0], [2.0]], device=device) + output = RankGaussian().fit_transform(TableTensor.from_tensor(context)) + probabilities = context.new_tensor([[0.25], [0.25], [0.75], [0.75]]) + torch.testing.assert_close( + output.numerical, torch.special.ndtri(probabilities) + ) + + +@withCUDA +@pytest.mark.parametrize("n_rows", [1, 4]) +def test_constants_and_missing_columns( + n_rows: int, device: torch.device +) -> None: + context = torch.tensor([[7.0, torch.nan]], device=device).expand( + n_rows, -1 + ) + processor = RankGaussian().fit(TableTensor.from_tensor(context)) + query = context.new_tensor( + [[7.0, 8.0], [torch.nan, torch.nan], [100.0, 3.0]] + ) + expected = context.new_tensor( + [[0.0, torch.nan], [torch.nan, torch.nan], [0.0, torch.nan]] + ) + torch.testing.assert_close( + processor.transform(TableTensor.from_tensor(query)).numerical, + expected, + equal_nan=True, + ) + + +@withCUDA +def test_nonfinite_context_does_not_change_observed_ranks( + device: torch.device, +) -> None: + context = torch.tensor( + [[0.0], [1.0], [torch.nan], [1.0], [torch.inf], [-torch.inf], [2.0]], + device=device, + ) + processor = RankGaussian().fit(TableTensor.from_tensor(context)) + reference = RankGaussian().fit( + TableTensor.from_tensor(context[[0, 1, 3, 6]]) + ) + query = TableTensor.from_tensor(context) + output = processor.transform(query).numerical + torch.testing.assert_close( + output, reference.transform(query).numerical, equal_nan=True + ) + assert torch.equal(output.isnan(), context.isnan()) + assert output[~context.isnan()].isfinite().all() + + +@withCUDA +def test_query_batch_does_not_change_fitted_ranks( + device: torch.device, +) -> None: + context = torch.arange(10.0, device=device).unsqueeze(-1) + processor = RankGaussian().fit(TableTensor.from_tensor(context)) + query = context.new_tensor([[1.5], [torch.nan], [5.5]]) + expected = processor.transform(TableTensor.from_tensor(query)).numerical + extended = torch.cat([query, context.new_tensor([[-1e12], [1e12]])]) + actual = processor.transform(TableTensor.from_tensor(extended)).numerical + torch.testing.assert_close(actual[:3], expected, equal_nan=True) + + +@withCUDA +@pytest.mark.parametrize("batch_shape", [(), (2,), (2, 3)]) +def test_batched_missing_values_match_independent_columns( + batch_shape: tuple[int, ...], device: torch.device +) -> None: + context = torch.randn(*batch_shape, 17, 4, device=device) + context[..., 1, 0] = torch.nan + context[..., :3, 1] = torch.nan + context[..., :, 2] = 2.0 + context[..., :, 3] = torch.nan + query = torch.randn(*batch_shape, 7, 4, device=device) + query[..., 0, 0] = torch.nan + processor = RankGaussian().fit(TableTensor.from_tensor(context)) + output = processor.transform(TableTensor.from_tensor(query)).numerical + for train, test, actual in zip( + context.reshape(-1, 17, 4), + query.reshape(-1, 7, 4), + output.reshape(-1, 7, 4), + strict=True, + ): + for column in range(4): + reference = RankGaussian().fit( + TableTensor.from_tensor(train[:, column : column + 1]) + ) + expected = reference.transform( + TableTensor.from_tensor(test[:, column : column + 1]) + ).numerical + torch.testing.assert_close( + actual[:, column : column + 1], expected, equal_nan=True + ) diff --git a/test/processing/test_contract.py b/test/processing/test_contract.py index 43692e47b..3646dbcf3 100644 --- a/test/processing/test_contract.py +++ b/test/processing/test_contract.py @@ -126,6 +126,7 @@ def _make_processor_pair( sp.QuantileTransform(n_quantiles=4, subsample=None), ), ProcessorCase(sp.Standardize()), + ProcessorCase(sp.RankGaussian()), ProcessorCase(sp.RobustScale()), ProcessorCase(sp.FlipSign()), ProcessorCase(sp.DropConstantColumns()), @@ -133,6 +134,7 @@ def _make_processor_pair( ProcessorCase(sp.RandomProjection(2)), ProcessorCase(sp.AlignCategories(), _make_align_categories_table), ProcessorCase(sp.ShuffleCategories()), + ProcessorCase(sp.ShuffleCategories(method="balanced_shift")), ProcessorCase(sp.ImputeMode()), ProcessorCase(sp.AddCategoryCounts()), ProcessorCase(sp.AddCalendarFields(["month"])), From 30943c74b7ec28fa74ee15df18f790cc7def2f83 Mon Sep 17 00:00:00 2001 From: Valter Hudovernik Date: Sat, 26 Sep 2026 04:14:25 -0700 Subject: [PATCH 2/3] Simplify Kumo transform schedule and shift dispatch --- sdm/models/kumo/tabular/recipe.py | 28 ++-------------------- sdm/processing/categorical/shuffle.py | 32 +++++++++++++------------ test/models/kumo/tabular/test_recipe.py | 26 ++++++++------------ 3 files changed, 29 insertions(+), 57 deletions(-) diff --git a/sdm/models/kumo/tabular/recipe.py b/sdm/models/kumo/tabular/recipe.py index 79c723099..9a22d4c59 100644 --- a/sdm/models/kumo/tabular/recipe.py +++ b/sdm/models/kumo/tabular/recipe.py @@ -22,6 +22,7 @@ def numerical_processor() -> sp.Sequential: sp.RobustScale(), sp.ClipSoft(3.0), ], + sp.RankGaussian(), method="round_robin", ), sp.ClipSigma(threshold=4.0), @@ -31,32 +32,7 @@ def numerical_processor() -> sp.Sequential: features=[ sp.StypeDispatch( numerical=[ - sp.Cast(torch.float64), - sp.DropConstantColumns(), - # Period 12 preserves the original three-way schedule - # while replacing every fourth view with Gaussian ranks. - sp.Choice( - *[ - [ - sp.RankGaussian(), - sp.Standardize(), - sp.ClipSigma(threshold=4.0), - ] - if i % 4 == 3 - else [ - sp.Standardize(eps=1e-6), - sp.Clip(-100.0, 100.0), - ( - sp.Identity(), - sp.PowerTransform(), - [sp.RobustScale(), sp.ClipSoft(3.0)], - )[i % 3], - sp.ClipSigma(threshold=4.0), - ] - for i in range(12) - ], - method="round_robin", - ), + numerical_processor(), sp.FlipSign(), ], categorical=[ diff --git a/sdm/processing/categorical/shuffle.py b/sdm/processing/categorical/shuffle.py index f9d6b0123..cc228abf1 100644 --- a/sdm/processing/categorical/shuffle.py +++ b/sdm/processing/categorical/shuffle.py @@ -63,25 +63,27 @@ def _draw_permutations( n_classes = category.numel() if n_classes <= 1: permutation = torch.arange(n_classes, device=device) - elif self.method in ("shift", "balanced_shift"): - if self.method == "balanced_shift": - key = (device, column, n_classes) - remaining = shifts.get(key) - if remaining is None or remaining.numel() == 0: - remaining = torch.randperm( - n_classes, - generator=generator, - device=device, - ) - offset = remaining[:1] - shifts[key] = remaining[1:] - else: - offset = torch.randint( + elif self.method == "shift": + offset = torch.randint( + n_classes, + (1,), + generator=generator, + device=device, + ) + permutation = ( + torch.arange(n_classes, device=device) - offset + ) % n_classes + elif self.method == "balanced_shift": + key = (device, column, n_classes) + remaining = shifts.get(key) + if remaining is None or remaining.numel() == 0: + remaining = torch.randperm( n_classes, - (1,), generator=generator, device=device, ) + offset = remaining[:1] + shifts[key] = remaining[1:] permutation = ( torch.arange(n_classes, device=device) - offset ) % n_classes diff --git a/test/models/kumo/tabular/test_recipe.py b/test/models/kumo/tabular/test_recipe.py index 4fdff2a07..75a0cfd80 100644 --- a/test/models/kumo/tabular/test_recipe.py +++ b/test/models/kumo/tabular/test_recipe.py @@ -169,23 +169,17 @@ def test_default_recipe_numerical_schedule( EnsembleTable.from_table(query, num_members=num_members) ) for i in range(num_members): - if i % 4 == 3: - processor = sp.Sequential( + processor = sp.Sequential( + sp.Standardize(eps=1e-6), + sp.Clip(-100.0, 100.0), + ( + sp.Identity(), + sp.PowerTransform(), + [sp.RobustScale(), sp.ClipSoft(3.0)], sp.RankGaussian(), - sp.Standardize(), - sp.ClipSigma(threshold=4.0), - ) - else: - processor = sp.Sequential( - sp.Standardize(eps=1e-6), - sp.Clip(-100.0, 100.0), - ( - sp.Identity(), - sp.PowerTransform(), - [sp.RobustScale(), sp.ClipSoft(3.0)], - )[i % 3], - sp.ClipSigma(threshold=4.0), - ) + )[i % 4], + sp.ClipSigma(threshold=4.0), + ) expected = processor.fit_transform(features).numerical.float() # Independent sign flips are allowed; magnitudes identify each view. torch.testing.assert_close( From f1bf2d233838172446d50f340493d93d9a7e3621 Mon Sep 17 00:00:00 2001 From: Valter Hudovernik Date: Sat, 26 Sep 2026 04:20:34 -0700 Subject: [PATCH 3/3] Remove redundant recipe schedule tests --- test/models/kumo/tabular/test_recipe.py | 64 ------------------------- 1 file changed, 64 deletions(-) diff --git a/test/models/kumo/tabular/test_recipe.py b/test/models/kumo/tabular/test_recipe.py index 75a0cfd80..a8a455719 100644 --- a/test/models/kumo/tabular/test_recipe.py +++ b/test/models/kumo/tabular/test_recipe.py @@ -150,67 +150,3 @@ def test_default_recipe_reduces_outputs_per_task() -> None: ) assert output.size() == (5, 3) torch.testing.assert_close(output.numerical.sum(dim=-1), torch.ones(5)) - - -@withCUDA -@pytest.mark.parametrize("num_members", [1, 4, 8, 16]) -def test_default_recipe_numerical_schedule( - num_members: int, device: torch.device -) -> None: - values = torch.arange(1, 101, device=device).double().square() - values[7] = torch.nan - features = TableTensor.from_tensor(values.unsqueeze(-1)) - recipe = KumoTabular.default_recipe() - output = recipe.features.fit_transform_ensemble( - EnsembleTable.from_table(features, num_members=num_members) - ) - query = features[:13] - prediction_features = recipe.features.transform_ensemble( - EnsembleTable.from_table(query, num_members=num_members) - ) - for i in range(num_members): - processor = sp.Sequential( - sp.Standardize(eps=1e-6), - sp.Clip(-100.0, 100.0), - ( - sp.Identity(), - sp.PowerTransform(), - [sp.RobustScale(), sp.ClipSoft(3.0)], - sp.RankGaussian(), - )[i % 4], - sp.ClipSigma(threshold=4.0), - ) - expected = processor.fit_transform(features).numerical.float() - # Independent sign flips are allowed; magnitudes identify each view. - torch.testing.assert_close( - output[i].numerical.abs(), expected.abs(), equal_nan=True - ) - torch.testing.assert_close( - prediction_features[i].numerical, - output[i].numerical[:13], - equal_nan=True, - ) - - -@withCUDA -@pytest.mark.parametrize("num_members", [1, 8, 16]) -@pytest.mark.parametrize("num_classes", [2, 3, 10]) -def test_default_recipe_balances_class_shifts( - num_members: int, num_classes: int, device: torch.device -) -> None: - target = TableTensor( - categorical=CategoricalTensor( - code=torch.arange(num_classes, device=device).unsqueeze(-1), - categories=(torch.arange(num_classes, device=device),), - ) - ) - output = KumoTabular.default_recipe().target.fit_transform_ensemble( - EnsembleTable.from_table(target, num_members=num_members) - ) - codes = torch.stack( - [output[i].categorical.code[0, 0] for i in range(num_members)] - ) - counts = codes.long().bincount(minlength=num_classes) - assert counts.max() - counts.min() <= 1 - for i in range(num_members): - assert output[i].categorical.tolist() == target.categorical.tolist()