diff --git a/changelog/1358.fixed.md b/changelog/1358.fixed.md new file mode 100644 index 000000000..5061fc691 --- /dev/null +++ b/changelog/1358.fixed.md @@ -0,0 +1 @@ +Fix silently wrong predictions on MPS with 65,536 or more training rows (torch <= 2.14). diff --git a/src/tabpfn/architectures/tabpfn_v2.py b/src/tabpfn/architectures/tabpfn_v2.py index 54b4775ee..5d23927d7 100644 --- a/src/tabpfn/architectures/tabpfn_v2.py +++ b/src/tabpfn/architectures/tabpfn_v2.py @@ -1259,11 +1259,12 @@ def _prepare_targets(y: torch.Tensor, num_rows: int, batch_size: int) -> torch.T if num_train_labels > num_rows: raise ValueError("No test rows provided.") target_RB1 = y.view(num_train_labels, 1 if y.ndim == 1 else batch_size, -1) - return torch.nn.functional.pad( - target_RB1, - (0, 0, 0, 0, 0, num_rows - num_train_labels), - value=float("nan"), + # Not F.pad: on MPS with torch <= 2.14, MPSGraph's constant pad corrupts the + # input once it reaches 2**16 rows, so every target came out NaN or misplaced. + nan_rows = target_RB1.new_full( + (num_rows - num_train_labels, *target_RB1.shape[1:]), float("nan") ) + return torch.cat([target_RB1, nan_rows]) def get_architecture(config: ArchitectureConfig) -> TabPFNV2: diff --git a/src/tabpfn/architectures/tabpfn_v2_5.py b/src/tabpfn/architectures/tabpfn_v2_5.py index b499d558a..4c7dff19f 100644 --- a/src/tabpfn/architectures/tabpfn_v2_5.py +++ b/src/tabpfn/architectures/tabpfn_v2_5.py @@ -1147,11 +1147,12 @@ def _prepare_targets( # Make sure the target is 3-dimensional. target_RBY = y.view(num_train_labels, 1 if y.ndim == 1 else batch_size, -1) - return torch.nn.functional.pad( - target_RBY, - (0, 0, 0, 0, 0, num_train_rows - num_train_labels), - value=float("nan"), + # Not F.pad: on MPS with torch <= 2.14, MPSGraph's constant pad corrupts the + # input once it reaches 2**16 rows, so every target came out NaN or misplaced. + nan_rows = target_RBY.new_full( + (num_train_rows - num_train_labels, *target_RBY.shape[1:]), float("nan") ) + return torch.cat([target_RBY, nan_rows]) def _replace_keys_from_base_architecture( diff --git a/src/tabpfn/architectures/tabpfn_v2_6.py b/src/tabpfn/architectures/tabpfn_v2_6.py index 0661d7695..6de94c435 100644 --- a/src/tabpfn/architectures/tabpfn_v2_6.py +++ b/src/tabpfn/architectures/tabpfn_v2_6.py @@ -1124,11 +1124,12 @@ def _prepare_targets( # Make sure the target is 3-dimensional. target_RBY = y.view(num_train_labels, 1 if y.ndim == 1 else batch_size, -1) - return torch.nn.functional.pad( - target_RBY, - (0, 0, 0, 0, 0, num_train_rows - num_train_labels), - value=float("nan"), + # Not F.pad: on MPS with torch <= 2.14, MPSGraph's constant pad corrupts the + # input once it reaches 2**16 rows, so every target came out NaN or misplaced. + nan_rows = target_RBY.new_full( + (num_train_rows - num_train_labels, *target_RBY.shape[1:]), float("nan") ) + return torch.cat([target_RBY, nan_rows]) def _remove_constant_features( diff --git a/src/tabpfn/architectures/tabpfn_v3.py b/src/tabpfn/architectures/tabpfn_v3.py index 0ce35f6a3..7d2fc30dc 100644 --- a/src/tabpfn/architectures/tabpfn_v3.py +++ b/src/tabpfn/architectures/tabpfn_v3.py @@ -2580,11 +2580,13 @@ def _prepare_targets( if num_train_labels > num_train_and_test_rows: raise ValueError("No test rows provided.") target_RBT = y.view(num_train_labels, 1 if y.ndim == 1 else batch_size, -1) - return F.pad( - target_RBT, - (0, 0, 0, 0, 0, num_train_and_test_rows - num_train_labels), - value=float("nan"), + # Not F.pad: on MPS with torch <= 2.14, MPSGraph's constant pad corrupts the + # input once it reaches 2**16 rows, so every target came out NaN or misplaced. + nan_rows = target_RBT.new_full( + (num_train_and_test_rows - num_train_labels, *target_RBT.shape[1:]), + float("nan"), ) + return torch.cat([target_RBT, nan_rows]) def _impute_nan_and_inf_with_mean( diff --git a/src/tabpfn/architectures/tabpfn_v3_5.py b/src/tabpfn/architectures/tabpfn_v3_5.py index 234d6dc00..37f7e4c36 100644 --- a/src/tabpfn/architectures/tabpfn_v3_5.py +++ b/src/tabpfn/architectures/tabpfn_v3_5.py @@ -3124,11 +3124,13 @@ def _prepare_targets( if num_train_labels > num_train_and_test_rows: raise ValueError("No test rows provided.") target_RBT = y.view(num_train_labels, 1 if y.ndim == 1 else batch_size, -1) - return F.pad( - target_RBT, - (0, 0, 0, 0, 0, num_train_and_test_rows - num_train_labels), - value=float("nan"), + # Not F.pad: on MPS with torch <= 2.14, MPSGraph's constant pad corrupts the + # input once it reaches 2**16 rows, so every target came out NaN or misplaced. + nan_rows = target_RBT.new_full( + (num_train_and_test_rows - num_train_labels, *target_RBT.shape[1:]), + float("nan"), ) + return torch.cat([target_RBT, nan_rows]) def _impute_nan_and_inf_with_mean( diff --git a/tests/test_architectures/test_prepare_targets.py b/tests/test_architectures/test_prepare_targets.py new file mode 100644 index 000000000..fefd4fbc6 --- /dev/null +++ b/tests/test_architectures/test_prepare_targets.py @@ -0,0 +1,38 @@ +"""Tests that `_prepare_targets` keeps the train targets intact at >= 2**16 rows. + +On MPS with torch <= 2.14, `F.pad` with a non-zero fill value overwrites its input +once it reaches 2**16 rows, which turned every train target into NaN or another +row's value. +""" + +from __future__ import annotations + +import pytest +import torch + +from tabpfn.architectures import ( + tabpfn_v2, + tabpfn_v2_5, + tabpfn_v2_6, + tabpfn_v3, + tabpfn_v3_5, +) + +_ARCHITECTURES = [tabpfn_v2, tabpfn_v2_5, tabpfn_v2_6, tabpfn_v3, tabpfn_v3_5] +_DEVICES = ["cpu"] + (["mps"] if torch.backends.mps.is_available() else []) + + +@pytest.mark.parametrize( + "architecture", _ARCHITECTURES, ids=lambda m: m.__name__.rsplit(".", 1)[-1] +) +@pytest.mark.parametrize("device", _DEVICES) +@pytest.mark.parametrize("num_train", [65535, 65536, 70000]) +@pytest.mark.parametrize("num_test", [0, 8]) +def test__prepare_targets__large_train_set__keeps_targets_and_nan_pads_test_rows( + architecture, device: str, num_train: int, num_test: int +) -> None: + y = torch.randn(num_train) + out = architecture._prepare_targets(y.to(device), num_train + num_test, 1).cpu() + assert out.shape == (num_train + num_test, 1, 1) + torch.testing.assert_close(out[:num_train, 0, 0], y, rtol=0, atol=0) + assert out[num_train:].isnan().all()