From 4474ae8d4ccb0509c789aa9218de5bed89245261 Mon Sep 17 00:00:00 2001
From: manjunathshiva <manjunath.shiva@gmail.com>
Date: Mon, 20 Jul 2026 09:48:53 +0530
Subject: [PATCH] Cast float64 targets before the device move in the
 context-cache path

_build_context_cache_pytorch moved y to the model device before casting
float64 down to float32. MPS rejects float64 tensors at transfer time, so
fit(cache_context=True) crashed on Apple Silicon for any float64 target
array (numpy's default float dtype). Same root cause as #68 / #69, which
cover the uncached _predict_step_pytorch path; this fixes the remaining
site. CPU/CUDA are unaffected: they accept float64 and the later cast
made the ordering unobservable there.

Adds a regression test that fits a small regressor with cache_context=True
on an MPS model with float64 targets; it fails with the TypeError against
the previous ordering and skips where MPS is unavailable.
---
 tabfm/src/classifier_and_regressor.py         |  3 +-
 .../classifier_and_regressor_pytorch_test.py  | 43 +++++++++++++++++++
 2 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/tabfm/src/classifier_and_regressor.py b/tabfm/src/classifier_and_regressor.py
index 4a0ebee..5d347fe 100644
--- a/tabfm/src/classifier_and_regressor.py
+++ b/tabfm/src/classifier_and_regressor.py
@@ -2064,9 +2064,10 @@ def _build_context_cache_pytorch(
         Xs_split, ys_split, cat_masks_split, ds_split
     ):
       X_t = torch.from_numpy(X_batch).to(device, dtype=torch.float32)
-      y_t = torch.from_numpy(y_batch).to(device)
+      y_t = torch.from_numpy(y_batch)
       if y_t.dtype == torch.float64:
         y_t = y_t.to(torch.float32)
+      y_t = y_t.to(device)
       cat_mask_t = torch.from_numpy(cat_mask_batch).to(device)
       d_t = torch.from_numpy(ds_batch).to(device)
       _, cache = model.prefill(X_t, y_t, cat_mask=cat_mask_t, d=d_t)
diff --git a/tabfm/src/classifier_and_regressor_pytorch_test.py b/tabfm/src/classifier_and_regressor_pytorch_test.py
index 4ea1df0..c31659d 100644
--- a/tabfm/src/classifier_and_regressor_pytorch_test.py
+++ b/tabfm/src/classifier_and_regressor_pytorch_test.py
@@ -136,6 +136,49 @@ def test_regressor_fit_predict(self):
     self.assertEqual(preds_cached.shape, (10,))
     np.testing.assert_allclose(preds_cached, preds, rtol=1e-5, atol=1e-6)
 
+  def test_regressor_cache_context_float64_targets_on_mps(self):
+    # _build_context_cache_pytorch moved y to the device before casting
+    # float64 down to float32. MPS rejects float64 tensors at transfer time,
+    # so fit(cache_context=True) crashed on Apple Silicon for any float64
+    # target array (numpy's default float dtype). CPU/CUDA accept float64,
+    # so the ordering is only observable on MPS.
+    if not torch.backends.mps.is_available():
+      self.skipTest("MPS is required for this test.")
+
+    np.random.seed(42)
+    model = pytorch_model.TabFM(
+        embed_dim=8,
+        max_classes=1,
+        col_num_blocks=1,
+        col_nhead=2,
+        col_num_inds=8,
+        row_num_blocks=1,
+        row_nhead=2,
+        row_num_cls=2,
+        icl_num_blocks=1,
+        icl_nhead=2,
+        ff_factor=2,
+        feature_group_size=2,
+        is_classifier=False
+    ).to("mps")
+
+    reg = TabFMRegressor(
+        model=model,
+        n_estimators=2,
+        batch_size=2,
+        random_state=42,
+        cache_context=True,
+        maybe_quantize_kv_cache=False,
+    )
+
+    X = np.random.rand(10, 3)
+    y = np.random.rand(10)  # float64 — numpy's default float dtype
+
+    reg.fit(X, y)
+    preds = reg.predict(X)
+    self.assertEqual(preds.shape, (10,))
+    self.assertTrue(np.all(np.isfinite(preds)))
+
 
 class PyTorchModelPickleTest(unittest.TestCase):
   """The PyTorch model must be picklable.