Skip to content

Commit 98c2948

Browse files
authored
Fix compiled refits with empty table schemas (#1241)
# Summary Fix compiled processor refits that produce an empty table schema. Preserve the selected column names and semantic types while making constructor metadata safe to emit in captured graphs. # Description of the change On PyTorch 2.14.0, Python 3.12.14, CPU/macOS, this public SDM workflow fails on the third refit: ```python import torch from sdm import TableTensor from sdm.processing import DropConstantColumns processor = DropConstantColumns() fit_transform = torch.compile( processor.fit_transform, backend="eager", fullgraph=False ) for data in ( torch.ones(3, 3), torch.arange(9, dtype=torch.float32).reshape(3, 3), torch.ones(3, 3), ): fit_transform(TableTensor.from_tensor(data)) ``` The generated constructor call contains enum keys such as `Stype.numerical` without binding `Stype`, raising `NameError: name 'Stype' is not defined`. Pass semantic type keys as their string values across the opaque table-constructor boundary, then restore the enum keys inside the constructor. This preserves the public schema and fixes the permissive compilation refit failure; strict fullgraph fitting with value-dependent column selection remains separate. # Validation - Table and constant-column suites: 63 passed, 14 skipped on CPU PyTorch 2.14.0. - External reproduction fails on unchanged main and passes with this change. - Repeated public refits with eager and Inductor backends preserve empty, nonempty, and zero-width output schemas. - Changed-file pre-commit checks passed, including formatting and type checking.
1 parent 05a88d9 commit 98c2948

1 file changed

Lines changed: 8 additions & 4 deletions

File tree

‎sdm/tensor/table.py‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -338,8 +338,9 @@ def __new__(
338338
if len(column_names) != len(set(column_names)):
339339
raise ValueError("Expected column names to be unique")
340340

341+
# Keep enum objects out of captured graph arguments.
341342
return TableTensor._make_table_tensor(
342-
columns=columns,
343+
columns={str(stype): names for stype, names in columns.items()},
343344
numerical=numerical,
344345
categorical=categorical,
345346
datetime=datetime,
@@ -352,7 +353,7 @@ def __new__(
352353
@staticmethod
353354
@torch.compiler.allow_in_graph
354355
def _make_table_tensor(
355-
columns: dict[Stype, tuple[str, ...]],
356+
columns: dict[str, tuple[str, ...]],
356357
numerical: Tensor,
357358
categorical: CategoricalTensor,
358359
datetime: Tensor,
@@ -361,9 +362,12 @@ def _make_table_tensor(
361362
cls: type[TableTensor] | None = None,
362363
) -> TableTensor:
363364
cls = TableTensor if cls is None else cls
365+
stype_columns = {
366+
Stype(stype): names for stype, names in columns.items()
367+
}
364368
column_to_loc = {
365369
name: (stype, i)
366-
for stype, names in columns.items()
370+
for stype, names in stype_columns.items()
367371
for i, name in enumerate(names)
368372
}
369373
out = Tensor._make_wrapper_subclass(
@@ -378,7 +382,7 @@ def _make_table_tensor(
378382
out._datetime = datetime
379383
out._text = text
380384
out._id = id
381-
out._columns = columns
385+
out._columns = stype_columns
382386
out._column_to_loc = column_to_loc
383387
return out
384388

0 commit comments

Comments
 (0)