Exclude NaN categories during alignment fitting

Author: aw471

sdm/processing/categorical/_categorical.py

@@ -2,10 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import os
-from collections.abc import Iterable, Sequence
-from typing import Any
 
-import torch
 from torch import Tensor
 
 from sdm import TableTensor
@@ -34,38 +31,6 @@ def _check_categorical_codes(table: TableTensor) -> None:
             )
 
 
-def _check_category_values(
-    categories: Iterable[Tensor],
-    columns: Sequence[str] | None = None,
-) -> None:
-    """Reject NaN values when fitting or loading categorical state."""
-    if os.getenv("SDM_ENABLE_CATEGORY_CHECKS", "1") == "0":
-        return
-    for index, values in enumerate(categories):
-        if not values.is_floating_point():
-            continue
-        if columns is None:
-            message = f"Fitted category vocabulary {index} contains NaN."
-        else:
-            message = (
-                f"Categorical column {columns[index]!r} contains NaN "
-                "in its category vocabulary."
-            )
-        valid = values.isnan().any().logical_not()
-        if torch.compiler.is_compiling():
-            torch._assert_async(valid.to("cpu"), message)
-        elif not valid:
-            raise ValueError(message)
-
-
-class _CategoryList(BufferList[Tensor]):
-    """Category buffers validated after loading fitted state."""
-
-    def _load_from_state_dict(self, *args: Any, **kwargs: Any) -> None:
-        super()._load_from_state_dict(*args, **kwargs)
-        _check_category_values(self)
-
-
 def _check_categories(
     table: TableTensor,
     categories: BufferList[Tensor],

sdm/processing/categorical/align.py

@@ -130,6 +130,8 @@ def _fit_column(
             ordered_categories = input_categories
 
         observed = counts.gather(1, order) >= self.min_frequency
+        if input_categories.is_floating_point():
+            observed &= ~input_categories.isnan()[order]
         # Only ragged vocabularies require per-batch materialization.
         fitted_categories = []
         for batch_index in range(batch_size):

sdm/processing/categorical/count.py

@@ -7,12 +7,11 @@
 from torch import Tensor
 
 from sdm import Stype, TableTensor
+from sdm.nn._buffer import BufferList
 from sdm.processing import Processor
 from sdm.processing.categorical._categorical import (
-    _CategoryList,
     _check_categorical_codes,
     _check_categories,
-    _check_category_values,
 )
 from sdm.tensor.categorical import _iter_categories
 
@@ -27,7 +26,6 @@ class AddCategoryCounts(Processor):
     Transform inputs must use the fitted per-column category vocabularies.
     Use :class:`~sdm.processing.AlignCategories` before this processor when
     training and transform inputs were tensorized independently.
-    Refit after changing category vocabularies; do not modify them in place.
 
     Args:
         min_cardinality: A count column is added when its fitted vocabulary
@@ -45,7 +43,7 @@ def __init__(self, min_cardinality: int = 0) -> None:
         super().__init__()
         self.min_cardinality = min_cardinality
         self._columns: tuple[str, ...] = ()
-        self._categories = _CategoryList()
+        self._categories: BufferList[Tensor] = BufferList()
         self.register_buffer("_log_counts", torch.empty(0))
         self.register_buffer(
             "_num_categories",
@@ -66,10 +64,6 @@ def _fit(
         *,
         generator: torch.Generator | None = None,
     ) -> None:
-        _check_category_values(
-            categories=_iter_categories(table.categorical),
-            columns=table.categorical_columns,
-        )
         _check_categorical_codes(table)
         categorical = table.categorical
         self._columns = tuple(
@@ -107,7 +101,7 @@ def _fit(
             value=counts.to(count_dtype).log1p().to(dtype),
         )
         self._set_buffer("_num_categories", num_categories)
-        self._categories = _CategoryList(_iter_categories(table.categorical))
+        self._categories = BufferList(_iter_categories(table.categorical))
 
     def _transform(self, table: TableTensor) -> TableTensor:
         _check_categories(table, self._categories)

sdm/processing/categorical/impute.py

@@ -4,14 +4,12 @@
 import torch
 
 from sdm import CategoricalTensor, Stype, TableTensor
+from sdm.nn._buffer import BufferList
 from sdm.processing import Processor
 from sdm.processing.categorical._categorical import (
-    _CategoryList,
     _check_categorical_codes,
     _check_categories,
-    _check_category_values,
 )
-from sdm.tensor.categorical import _iter_categories
 
 
 class ImputeMode(Processor):
@@ -25,7 +23,6 @@ class ImputeMode(Processor):
     Transform inputs must use the fitted per-column category vocabularies.
     Use :class:`~sdm.processing.AlignCategories` before this processor when
     training and transform inputs were tensorized independently.
-    Refit after changing category vocabularies; do not modify them in place.
     """
 
     handles_stypes = (Stype.categorical,)
@@ -37,7 +34,7 @@ def __init__(self) -> None:
         super().__init__()
         # TODO: Add a separate processor that encodes missing values as their
         # own category instead of imputing an observed one.
-        self._categories = _CategoryList()
+        self._categories: BufferList[torch.Tensor] = BufferList()
         self.register_buffer(
             "_fill_values",
             torch.empty(0, dtype=torch.long),
@@ -49,10 +46,6 @@ def _fit(
         *,
         generator: torch.Generator | None = None,
     ) -> None:
-        _check_category_values(
-            categories=_iter_categories(table.categorical),
-            columns=table.categorical_columns,
-        )
         data = table.categorical
         _check_categorical_codes(table)
         fill_values: list[torch.Tensor] = []
@@ -86,7 +79,7 @@ def _fit(
             if len(fill_values) > 0
             else torch.empty(0, dtype=torch.long, device=data.device),
         )
-        self._categories = _CategoryList(table.categorical.categories)
+        self._categories = BufferList(table.categorical.categories)
 
     def _fit_transform(
         self,

test/processing/categorical/test_align.py

@@ -295,8 +295,10 @@ def test_align_categories_numeric_values(device: torch.device) -> None:
 
 
 @withCUDA
+@pytest.mark.parametrize("sort_by", ["code", "frequency", "value"])
 def test_align_categories_does_not_match_nan_category_values(
     device: torch.device,
+    sort_by: Literal["code", "frequency", "value"],
 ) -> None:
     context = TableTensor(
         categorical=CategoricalTensor(
@@ -315,9 +317,13 @@ def test_align_categories_does_not_match_nan_category_values(
         ),
     )
 
-    output = AlignCategories().fit(context).transform(query)
+    processor = AlignCategories(sort_by=sort_by)
+    aligned = processor.fit_transform(context)
+    output = processor.transform(query)
 
-    assert output.categorical.code.squeeze(-1).tolist() == [-1, -1, -1]
+    assert aligned.categorical.categories[0].tolist() == [1.0]
+    assert aligned.categorical.code.squeeze(-1).tolist() == [-1, 0]
+    assert output.categorical.code.squeeze(-1).tolist() == [-1, -1, 0]
 
 
 def test_align_categories_scales_to_large_numeric_vocabulary() -> None:

test/processing/categorical/test_impute.py

@@ -6,30 +6,9 @@
 
 import sdm.processing as sp
 from sdm import CategoricalTensor, StringTensor, TableTensor
-from sdm.processing import Processor
 from sdm.testing import withCUDA
 
 
-@withCUDA
-@pytest.mark.parametrize(
-    "processor_type", [sp.AddCategoryCounts, sp.ImputeMode]
-)
-def test_float_category_state_roundtrip(
-    device: torch.device,
-    processor_type: type[Processor],
-) -> None:
-    table = TableTensor(
-        categorical=CategoricalTensor(
-            code=torch.tensor([[0], [1], [0], [-1]], device=device),
-            categories=(torch.tensor([1.0, 2.0], device=device),),
-        ),
-    )
-    processor = processor_type().fit(table)
-    restored = processor_type().to(device=device)
-    restored.load_state_dict(processor.state_dict())
-    assert restored.transform(table).equal(processor.transform(table))
-
-
 def _table(
     values: list[list[int]] | list[list[list[int]]],
     *,