Fix silently wrong targets on MPS with >= 65,536 training rows

changelog/1358.fixed.md

@@ -0,0 +1 @@
+Fix silently wrong predictions on MPS with 65,536 or more training rows (torch <= 2.14).

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:

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(

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(

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(

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(

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()