fix(inference): validate tuning holdout frac before inverting it for folds

changelog/1350.fixed.md

@@ -0,0 +1 @@
+Validate 	uning_holdout_frac and derive the tuning fold count from the requested fraction, so an unsupported value reports a clear error instead of an unusable fold count.

src/tabpfn/inference_tuning.py

@@ -6,6 +6,7 @@
 
 import copy
 import dataclasses
+import math
 import warnings
 from collections.abc import Callable
 from enum import Enum
@@ -34,6 +35,15 @@
 
 MIN_NUM_SAMPLES_RECOMMENDED_FOR_TUNING = 500
 
+MAX_TUNING_FOLDS = 100
+"""Upper bound on the (stratified) k-fold count implied by a holdout fraction.
+
+A smaller holdout fraction means more folds, and folds are what make each tuning
+split expensive, so we refuse fractions that would ask for more than this many.
+"""
+MIN_TUNING_HOLDOUT_FRAC = 1 / MAX_TUNING_FOLDS
+"""The smallest holdout fraction that stays within `MAX_TUNING_FOLDS`."""
+
 
 @dataclasses.dataclass
 class TuningConfig:
@@ -235,6 +245,37 @@ def compute_regression_metric_to_minimize(
     )
 
 
+def _folds_for_holdout_frac(holdout_frac: float) -> int:
+    """Derives the number of (stratified) k-folds from a holdout fraction.
+
+    Each k-fold split holds out `1 / n_folds` of the data, so the requested
+    fraction has to be inverted to get the fold count.
+
+    Args:
+        holdout_frac: The fraction of the data to hold out for tuning.
+
+    Returns:
+        The number of folds to use, in `[2, MAX_TUNING_FOLDS]`.
+
+    Raises:
+        ValueError: If `holdout_frac` is not a finite number in the open
+            interval `(0, 1)`, or if it would need more than
+            `MAX_TUNING_FOLDS` folds.
+    """
+    if not math.isfinite(holdout_frac) or not 0.0 < holdout_frac < 1.0:
+        raise ValueError(
+            "tuning holdout fraction must be a finite number strictly between "
+            f"0 and 1, got {holdout_frac!r}."
+        )
+    if holdout_frac < MIN_TUNING_HOLDOUT_FRAC:
+        raise ValueError(
+            f"tuning holdout fraction {holdout_frac!r} is too small: it would "
+            f"need more than {MAX_TUNING_FOLDS} tuning folds. Use a value of at "
+            f"least {MIN_TUNING_HOLDOUT_FRAC}."
+        )
+    return min(max(round(1 / holdout_frac), 2), MAX_TUNING_FOLDS)
+
+
 def get_tuning_splits(
     X: np.ndarray,
     y: np.ndarray,
@@ -260,13 +301,15 @@ def get_tuning_splits(
         Returns a list of splits as tuples of
         (X_train_NtF, X_holdout_NhF, y_train_Nt, y_holdout_Nh).
         Shape suffixes: Nt=num train samples, F=num features, Nh=num holdout samples.
+
+    Raises:
+        ValueError: If `holdout_frac` is not a finite number in the open
+            interval `(0, 1)`.
     """
     # We want to use (Stratified)KFold to ensure that no train samples are used twice.
     # Therefore, we have to invert the holdout_frac to get the number of folds to
-    # use for (Stratified)KFold. Round holdout_frac to 2 digits to avoid needing
-    # more than 100 folds
-    rounded_holdout_frac = round(holdout_frac, 2)
-    n_folds = max(2, round(1 / rounded_holdout_frac))
+    # use for (Stratified)KFold.
+    n_folds = _folds_for_holdout_frac(holdout_frac)
 
     if isinstance(random_state, np.random.Generator):
         # (Stratified)KFold does not accept np.random.Generator.

tests/test_inference_tuning.py

@@ -560,6 +560,108 @@ def test__get_tuning_splits__classification_is_unchanged_by_default() -> None:
         assert (y_holdout == 0).sum() == 4
 
 
+@pytest.mark.parametrize(
+    "holdout_frac",
+    [
+        pytest.param(0.0, id="zero"),
+        pytest.param(-0.1, id="negative"),
+        pytest.param(1.0, id="one"),
+        pytest.param(1.5, id="above_one"),
+        pytest.param(float("nan"), id="nan"),
+        pytest.param(float("inf"), id="inf"),
+    ],
+)
+def test__get_tuning_splits__rejects_holdout_frac_outside_open_unit_interval(
+    holdout_frac: float,
+) -> None:
+    y = np.random.default_rng(0).normal(size=200)
+
+    # A frac that rounds to zero previously raised ZeroDivisionError, and NaN or
+    # out-of-range fracs silently produced a nonsense fold count.
+    with pytest.raises(ValueError, match="strictly between"):
+        get_tuning_splits(
+            X=_identity_X(200),
+            y=y,
+            holdout_frac=holdout_frac,
+            n_splits=1,
+            task_type="regressor",
+        )
+
+
+@pytest.mark.parametrize(
+    ("holdout_frac", "expected_holdout_size"),
+    [
+        pytest.param(0.15, 143, id="fractional_frac"),
+        pytest.param(0.01, 10, id="at_the_supported_lower_bound"),
+    ],
+)
+def test__get_tuning_splits__derives_sane_fold_count_from_awkward_frac(
+    holdout_frac: float,
+    expected_holdout_size: int,
+) -> None:
+    n_samples = 1_000
+    y = np.random.default_rng(0).normal(size=n_samples)
+
+    splits = get_tuning_splits(
+        X=_identity_X(n_samples),
+        y=y,
+        holdout_frac=holdout_frac,
+        n_splits=1,
+        task_type="regressor",
+    )
+
+    # The fold count is derived from the exact frac, so 0.15 maps to the nearest
+    # usable fold count (7) instead of the 6 that pre-rounding implied.
+    assert len(splits[0][1]) == expected_holdout_size
+
+
+@pytest.mark.parametrize("holdout_frac", [0.004, 0.005, 0.009])
+def test__get_tuning_splits__rejects_frac_needing_more_folds_than_supported(
+    holdout_frac: float,
+) -> None:
+    y = np.random.default_rng(0).normal(size=200)
+
+    # These all need more than the supported fold count. Previously 0.004 divided
+    # by zero outright, and 0.005 asked for 100 folds, which StratifiedKFold
+    # rejects on most real datasets.
+    with pytest.raises(ValueError, match="too small"):
+        get_tuning_splits(
+            X=_identity_X(200),
+            y=y,
+            holdout_frac=holdout_frac,
+            n_splits=1,
+            task_type="regressor",
+        )
+
+
+@pytest.mark.parametrize(
+    ("holdout_frac", "expected_holdout_size"),
+    [
+        pytest.param(0.1, 100, id="auto_default_small_dataset"),
+        pytest.param(0.2, 200, id="auto_default_mid_dataset"),
+        pytest.param(0.3, 334, id="auto_default_large_dataset"),
+    ],
+)
+def test__get_tuning_splits__auto_default_fracs_keep_their_fold_counts(
+    holdout_frac: float,
+    expected_holdout_size: int,
+) -> None:
+    # `get_default_tuning_holdout_frac` only ever yields 0.1/0.2/0.3. Those must
+    # keep implying 10/5/3 folds, so existing users see no behaviour change.
+    n_samples = 1_000
+    y = np.random.default_rng(0).normal(size=n_samples)
+
+    splits = get_tuning_splits(
+        X=_identity_X(n_samples),
+        y=y,
+        holdout_frac=holdout_frac,
+        n_splits=1,
+        task_type="regressor",
+    )
+
+    assert len(splits[0][1]) == expected_holdout_size
+
+
 @pytest.mark.parametrize(
     ("num_samples", "expected_holdout_frac", "expected_n_folds"),
     [