fix(classifier, regressor): build the feature schema from the column count

changelog/1349.fixed.md

@@ -0,0 +1 @@
+Build the categorical feature schema from the number of features rather than the number of training rows, so a categorical column at a high index is reported to the model as categorical.

src/tabpfn/classifier.py

@@ -1082,9 +1082,11 @@ def fit_from_preprocessed(
             self.n_classes_ = max(int(t.max().item()) for t in y_preprocessed) + 1
             self.classes_ = torch.arange(self.n_classes_)
 
+        # X_preprocessed is the collated batch: (n_datasets, n_train_rows,
+        # n_columns), so the column count is the last dimension.
         feature_schema = convert_batch_of_cat_ix_to_schema(
             batch_of_cat_indices=cat_ix,
-            num_features=X_preprocessed[0].shape[1],
+            num_features=X_preprocessed[0].shape[-1],
         )
 
         # Preprocessed tensors hold no dates or strings either, so these fit

src/tabpfn/regressor.py

@@ -1147,9 +1147,11 @@ def fit_from_preprocessed(
                 self.inference_precision, self.devices_
             )
 
+        # X_preprocessed is the collated batch: (n_datasets, n_train_rows,
+        # n_columns), so the column count is the last dimension.
         feature_schema = convert_batch_of_cat_ix_to_schema(
             batch_of_cat_indices=cat_ix,
-            num_features=X_preprocessed[0].shape[1],
+            num_features=X_preprocessed[0].shape[-1],
         )
 
         # Preprocessed tensors hold no dates or strings either, so these fit

tests/test_classifier_interface.py

@@ -24,6 +24,7 @@
 
 from tabpfn import TabPFNClassifier
 from tabpfn.architectures import tabpfn_v2_5
+from tabpfn.architectures.interface import PerformanceOptions
 from tabpfn.base import ModelSpecs, initialize_tabpfn_model
 from tabpfn.constants import ModelVersion
 from tabpfn.inference_config import DEFAULT_SOFTMAX_TEMPERATURE, InferenceConfig
@@ -34,6 +35,10 @@
 )
 from tabpfn.model_loading import ModelSource, prepend_cache_path
 from tabpfn.preprocessing import PreprocessorConfig
+from tabpfn.preprocessing.datamodel import FeatureModality
+from tabpfn.preprocessing.ensemble import (
+    generate_classification_ensemble_configs,
+)
 from tabpfn.utils import infer_devices
 
 from .utils import (
@@ -1742,3 +1747,142 @@ def mk(seed: int, n: int = 60, f: int = 5) -> tuple[np.ndarray, np.ndarray]:
     fitted.predict_proba_batched(X_list, y_list, X_tests)
     after = fitted.predict_proba(a_x[:5])
     np.testing.assert_array_equal(before, after)
+
+
+# =============================================================================
+# fit_from_preprocessed: the feature schema is built from the column count
+# =============================================================================
+
+
+N_BATCH_ROWS = 5
+N_BATCH_COLS = 8
+BATCH_CAT_COL = N_BATCH_COLS - 1  # index 7 >= the 5 rows of a collated batch
+
+
+def _collated_batch(
+    n_datasets: int = 2, n_estimators: int = 2
+) -> tuple[
+    ModelSpecs,
+    list[torch.Tensor],
+    list[torch.Tensor],
+    list[list[list[int]]],
+    list[list[object]],
+]:
+    """The (n_datasets, n_train_rows, n_columns) tensors the collator produces."""
+    specs = _create_dummy_classifier_model_specs()
+    flat_configs = generate_classification_ensemble_configs(
+        num_estimators=n_estimators,
+        add_fingerprint_feature=False,
+        polynomial_features="no",
+        feature_shift_decoder="shuffle",
+        preprocessor_configs=[
+            PreprocessorConfig(
+                "none", categorical_name="numeric", max_features_per_estimator=500
+            )
+        ],
+        class_shift_method="shuffle",
+        n_classes=2,
+        random_state=0,
+        num_models=1,
+        outlier_removal_std=None,
+    )
+    configs: list[list[object]] = [
+        [flat_configs[estimator] for estimator in range(n_estimators)]
+        for _ in range(n_datasets)
+    ]
+    torch.manual_seed(0)
+    X_preprocessed = [
+        torch.randn(n_datasets, N_BATCH_ROWS, N_BATCH_COLS),
+    ]
+    y_preprocessed = [
+        torch.randint(0, 2, (n_datasets, N_BATCH_ROWS)).float(),
+    ]
+    cat_ix = [[[BATCH_CAT_COL] for _ in range(n_estimators)] for _ in range(n_datasets)]
+    return specs, X_preprocessed, y_preprocessed, cat_ix, configs
+
+
+def test__fit_from_preprocessed__feature_schema_spans_the_columns() -> None:
+    """The schema is as wide as the batch's columns, not its rows.
+
+    X_preprocessed[0] is the collated (n_datasets, n_train_rows, n_columns)
+    tensor, so `shape[1]` is the row count. A schema built from it is narrower
+    than the data whenever there are more columns than rows, which silently
+    drops every categorical column whose index is >= n_train_rows and tells the
+    model those columns are numerical.
+    """
+    specs, X_preprocessed, y_preprocessed, cat_ix, configs = _collated_batch()
+
+    clf = TabPFNClassifier(
+        model_path=specs, device="cpu", n_estimators=2, random_state=0
+    )
+    clf.fit_from_preprocessed(
+        X_preprocessed,
+        y_preprocessed,
+        cat_ix,
+        configs,
+        performance_options=PerformanceOptions(),
+    )
+
+    assert X_preprocessed[0].shape[1] == N_BATCH_ROWS
+    assert X_preprocessed[0].shape[-1] == N_BATCH_COLS
+    for dataset_schemas in clf.executor_.feature_schema_list:
+        for schema in dataset_schemas:
+            assert len(schema.features) == N_BATCH_COLS
+            assert schema.indices_for(FeatureModality.CATEGORICAL) == [BATCH_CAT_COL]
+
+
+def test__predict_proba_batched__tells_the_model_about_high_categorical_columns() -> (
+    None
+):
+    """Batched inference reports the same categorical columns as the single path.
+
+    A wide dataset pushes the categorical past the row count after
+    preprocessing (the SVD/fingerprint steps add columns), so a row-wide schema
+    omits it and the model is told the column is numerical -- while
+    `predict_proba`, which uses each member's own schema, still reports it.
+    """
+    n_rows = 100  # MIN_NUMBER_OF_SAMPLES_FOR_CATEGORICAL_INFERENCE
+    n_cols = 250  # wide enough that the categorical lands beyond n_train_rows
+    rng = np.random.RandomState(0)
+    X = rng.randn(n_rows, n_cols).astype(object)
+    X[:, -1] = rng.choice(["a", "b", "c", "d"], size=n_rows)
+    y = (X[:, -1] == "a").astype(int)
+    X_test = np.repeat(X[:1], 2, axis=0).astype(object)
+    X_test[:, -1] = ["a", "b"]
+
+    clf = TabPFNClassifier(
+        model_path=_create_dummy_classifier_model_specs(),
+        device="cpu",
+        n_estimators=1,
+        random_state=0,
+    )
+    clf.fit(X, y)
+
+    member_cats = clf.executor_.ensemble_members[0].feature_schema.indices_for(
+        FeatureModality.CATEGORICAL
+    )
+    # Precondition: this is the case the bug needs, so the test cannot pass
+    # vacuously on a schema that happened to keep the column.
+    assert [c for c in member_cats if c >= n_rows], (
+        f"expected a categorical column at index >= {n_rows}, got {member_cats}"
+    )
+
+    seen: list[object] = []
+    original_forward = type(clf.models_[0]).forward
+
+    def spy(self: object, x: torch.Tensor, y: torch.Tensor, **kwargs: object) -> object:
+        seen.append(kwargs.get("categorical_inds"))
+        return original_forward(self, x, y, **kwargs)  # type: ignore[arg-type]
+
+    type(clf.models_[0]).forward = spy  # type: ignore[method-assign]
+    try:
+        clf.predict_proba(X_test)
+        single = list(seen)
+        seen.clear()
+        clf.predict_proba_batched([X], [y], [X_test])
+        batched = list(seen)
+    finally:
+        type(clf.models_[0]).forward = original_forward  # type: ignore[method-assign]
+
+    assert single, "the spy recorded no forward calls for predict_proba"
+    assert batched == single

tests/test_regressor_interface.py

@@ -26,6 +26,7 @@
 import tabpfn.regressor as regressor_module
 from tabpfn import TabPFNRegressor
 from tabpfn.architectures import tabpfn_v2_5
+from tabpfn.architectures.interface import PerformanceOptions
 from tabpfn.architectures.shared.bar_distribution import FullSupportBarDistribution
 from tabpfn.base import ModelSpecs, initialize_tabpfn_model
 from tabpfn.constants import ModelVersion
@@ -39,6 +40,8 @@
 )
 from tabpfn.model_loading import ModelSource, prepend_cache_path
 from tabpfn.preprocessing import PreprocessorConfig
+from tabpfn.preprocessing.datamodel import FeatureModality
+from tabpfn.preprocessing.ensemble import generate_regression_ensemble_configs
 from tabpfn.settings import settings
 from tabpfn.utils import infer_devices
 from tabpfn.validation import ensure_compatible_predict_input_sklearn
@@ -2100,3 +2103,139 @@ def test__eval_metric_and_tuning_config__round_trip_through_sklearn_clone(
     round_tripped = TabPFNRegressor().set_params(**params)
     assert round_tripped.eval_metric == eval_metric
     assert round_tripped.tuning_config == tuning_config
+
+
+# =============================================================================
+# fit_from_preprocessed: the feature schema is built from the column count
+# =============================================================================
+
+
+N_BATCH_ROWS = 5
+N_BATCH_COLS = 8
+BATCH_CAT_COL = N_BATCH_COLS - 1  # index 7 >= the 5 rows of a collated batch
+
+
+def _collated_batch(
+    n_datasets: int = 2, n_estimators: int = 2
+) -> tuple[
+    ModelSpecs,
+    list[torch.Tensor],
+    list[torch.Tensor],
+    list[list[list[int]]],
+    list[list[object]],
+]:
+    """The (n_datasets, n_train_rows, n_columns) tensors the collator produces."""
+    specs = _create_dummy_regressor_model_specs()
+    flat_configs = generate_regression_ensemble_configs(
+        num_estimators=n_estimators,
+        add_fingerprint_feature=False,
+        polynomial_features="no",
+        feature_shift_decoder="shuffle",
+        preprocessor_configs=[
+            PreprocessorConfig(
+                "none", categorical_name="numeric", max_features_per_estimator=500
+            )
+        ],
+        target_transforms=[None],
+        random_state=0,
+        num_models=1,
+        outlier_removal_std=None,
+    )
+    configs: list[list[object]] = [
+        [flat_configs[estimator] for estimator in range(n_estimators)]
+        for _ in range(n_datasets)
+    ]
+    torch.manual_seed(0)
+    X_preprocessed = [
+        torch.randn(n_datasets, N_BATCH_ROWS, N_BATCH_COLS),
+    ]
+    y_preprocessed = [
+        torch.randn(n_datasets, N_BATCH_ROWS),
+    ]
+    cat_ix = [[[BATCH_CAT_COL] for _ in range(n_estimators)] for _ in range(n_datasets)]
+    return specs, X_preprocessed, y_preprocessed, cat_ix, configs
+
+
+def test__fit_from_preprocessed__feature_schema_spans_the_columns() -> None:
+    """The schema is as wide as the batch's columns, not its rows.
+
+    X_preprocessed[0] is the collated (n_datasets, n_train_rows, n_columns)
+    tensor, so `shape[1]` is the row count. A schema built from it is narrower
+    than the data whenever there are more columns than rows, which silently
+    drops every categorical column whose index is >= n_train_rows and tells the
+    model those columns are numerical.
+    """
+    specs, X_preprocessed, y_preprocessed, cat_ix, configs = _collated_batch()
+
+    reg = TabPFNRegressor(
+        model_path=specs, device="cpu", n_estimators=2, random_state=0
+    )
+    reg.fit_from_preprocessed(
+        X_preprocessed,
+        y_preprocessed,
+        cat_ix,
+        configs,
+        performance_options=PerformanceOptions(),
+    )
+
+    assert X_preprocessed[0].shape[1] == N_BATCH_ROWS
+    assert X_preprocessed[0].shape[-1] == N_BATCH_COLS
+    for dataset_schemas in reg.executor_.feature_schema_list:
+        for schema in dataset_schemas:
+            assert len(schema.features) == N_BATCH_COLS
+            assert schema.indices_for(FeatureModality.CATEGORICAL) == [BATCH_CAT_COL]
+
+
+def test__predict_batched__tells_the_model_about_high_categorical_columns() -> None:
+    """Batched inference reports the same categorical columns as the single path.
+
+    A wide dataset pushes the categorical past the row count after
+    preprocessing (the SVD/fingerprint steps add columns), so a row-wide schema
+    omits it and the model is told the column is numerical -- while `predict`,
+    which uses each member's own schema, still reports it.
+    """
+    n_rows = 100  # wide enough to be a realistic regression context
+    n_cols = 250  # wide enough that the categorical lands beyond n_train_rows
+    rng = np.random.RandomState(0)
+    X = rng.randn(n_rows, n_cols).astype(object)
+    X[:, -1] = rng.choice(["a", "b", "c", "d"], size=n_rows)
+    y = rng.randn(n_rows)
+    X_test = np.repeat(X[:1], 2, axis=0).astype(object)
+    X_test[:, -1] = ["a", "b"]
+
+    reg = TabPFNRegressor(
+        model_path=_create_dummy_regressor_model_specs(),
+        device="cpu",
+        n_estimators=1,
+        random_state=0,
+    )
+    reg.fit(X, y)
+
+    member_cats = reg.executor_.ensemble_members[0].feature_schema.indices_for(
+        FeatureModality.CATEGORICAL
+    )
+    # Precondition: this is the case the bug needs, so the test cannot pass
+    # vacuously on a schema that happened to keep the column.
+    assert [c for c in member_cats if c >= n_rows], (
+        f"expected a categorical column at index >= {n_rows}, got {member_cats}"
+    )
+
+    seen: list[object] = []
+    original_forward = type(reg.models_[0]).forward
+
+    def spy(self: object, x: torch.Tensor, y: torch.Tensor, **kwargs: object) -> object:
+        seen.append(kwargs.get("categorical_inds"))
+        return original_forward(self, x, y, **kwargs)  # type: ignore[arg-type]
+
+    type(reg.models_[0]).forward = spy  # type: ignore[method-assign]
+    try:
+        reg.predict(X_test)
+        single = list(seen)
+        seen.clear()
+        reg.predict_batched([X], [y], [X_test])
+        batched = list(seen)
+    finally:
+        type(reg.models_[0]).forward = original_forward  # type: ignore[method-assign]
+
+    assert single, "the spy recorded no forward calls for predict"
+    assert batched == single