Add MLX backend for Apple silicon

.github/workflows/pytest_and_autopublish.yml

@@ -38,7 +38,7 @@ jobs:
         # cache-dependency-path: '**/pyproject.toml'
 
     - run: pip --version
-    - run: pip install -e .[dev,jax,pytorch]
+    - run: pip install -e .[dev,jax,pytorch,mlx]
     - run: pip freeze
 
     # Run tests (in parallel)

CHANGELOG.md

@@ -21,8 +21,19 @@ To release a new version (e.g. from `1.0.0` -> `2.0.0`):
 
 -->
 
+## [Unreleased]
+
+* Added an MLX backend (`tabfm/src/mlx/`, `pip install -e .[mlx]`) for native
+  Apple-silicon inference. It reuses the PyTorch v1.0.0 weight release
+  (identical parameter names/layouts) and is parity-tested against the PyTorch
+  port to < 1e-4 max abs diff in float32.
+* Fixed the `pytorch` extra missing `safetensors`: with a bare `torch`
+  install, `tabfm_v1_0_0_pytorch.load()` raised `NameError` inside
+  `PyTorchModelHubMixin` when loading the safetensors release.
+
 ## [1.0.0] - 2026-06-29
 
 * Initial release
 
+[Unreleased]: https://github.com/google-research/tabfm/compare/v1.0.0...HEAD
 [1.0.0]: https://github.com/google-research/tabfm/releases/tag/v1.0.0

README.md

@@ -34,6 +34,14 @@ pip install -e .[pytorch]
 ```
 *Note: For PyTorch with GPU support, ensure you have the appropriate PyTorch version installed for your CUDA version before installing TabFM.*
 
+**MLX (Apple silicon):**
+```bash
+git clone https://github.com/google-research/tabfm.git
+cd tabfm
+pip install -e .[mlx]
+```
+*Note: The MLX backend runs natively on Apple silicon (Metal / unified memory) and loads the PyTorch v1.0.0 weight release directly — no separate checkpoint is needed.*
+
 ### Requirements
 For a complete list of pinned dependencies and versions, please see [requirements.txt](requirements.txt). The core requirements depend on the backend you choose:
 *   Python >= 3.11
@@ -43,12 +51,14 @@ For a complete list of pinned dependencies and versions, please see [requirement
     *   Flax (specifically `flax==0.12.7`, using the modern `flax.nnx` API)
 *   **PyTorch Backend:**
     *   PyTorch (specifically `torch==2.12.1+cpu` or a GPU version)
+*   **MLX Backend:**
+    *   MLX (`mlx>=0.31`; Apple silicon recommended)
 
 ---
 
 ## Quick Start (TabFM v1.0.0)
 
-We provide pre-trained weights for the **TabFM v1.0.0** release. The library handles downloading and loading these weights automatically. You can choose to load the model using either the JAX or PyTorch backend.
+We provide pre-trained weights for the **TabFM v1.0.0** release. The library handles downloading and loading these weights automatically. You can choose to load the model using the JAX, PyTorch, or MLX backend.
 
 ### 1. Classification Example
 
@@ -67,7 +77,11 @@ model = tabfm_v1_0_0.load()
 # from tabfm import tabfm_v1_0_0_pytorch as tabfm_v1_0_0
 # model = tabfm_v1_0_0.load()
 
-# Initialize scikit-learn compatible classifier (works with either backend model)
+# OPTION C: MLX Backend (Apple silicon)
+# from tabfm import tabfm_v1_0_0_mlx as tabfm_v1_0_0
+# model = tabfm_v1_0_0.load()
+
+# Initialize scikit-learn compatible classifier (works with any backend model)
 clf = TabFMClassifier(model=model)
 
 # Prepare your dataset (supports mixed numerical and categorical features)
@@ -112,7 +126,11 @@ model = tabfm_v1_0_0.load(model_type="regression")
 # from tabfm import tabfm_v1_0_0_pytorch as tabfm_v1_0_0
 # model = tabfm_v1_0_0.load(model_type="regression")
 
-# Initialize scikit-learn compatible regressor (works with either backend model)
+# OPTION C: MLX Backend (Apple silicon)
+# from tabfm import tabfm_v1_0_0_mlx as tabfm_v1_0_0
+# model = tabfm_v1_0_0.load(model_type="regression")
+
+# Initialize scikit-learn compatible regressor (works with any backend model)
 reg = TabFMRegressor(model=model)
 
 # Prepare your dataset
@@ -167,6 +185,10 @@ PYTHONPATH=. python3 -m unittest discover -s tabfm/src/ -p "*_test.py"
 # Or run specific test files:
 PYTHONPATH=. python3 -m unittest tabfm/src/pytorch/model_test.py
 PYTHONPATH=. python3 -m unittest tabfm/src/classifier_and_regressor_pytorch_test.py
+
+# MLX backend tests (require mlx; the parity test also requires torch):
+PYTHONPATH=. python3 -m unittest tabfm/src/mlx/model_test.py
+PYTHONPATH=. python3 -m unittest tabfm/src/classifier_and_regressor_mlx_test.py
 ```
 
 Alternatively, if you have Bazel installed, you can run tests with:

conftest.py

@@ -33,6 +33,7 @@
 # backend is not installed.
 _has_torch = importlib.util.find_spec("torch") is not None
 _has_jax = importlib.util.find_spec("jax") is not None
+_has_mlx = importlib.util.find_spec("mlx") is not None
 collect_ignore = []
 if not _has_torch:
   collect_ignore.append("tabfm/src/classifier_and_regressor_pytorch_test.py")
@@ -42,10 +43,15 @@
       "tabfm/src/jax/checkpointing_test.py",
       "tabfm/src/jax/memory_efficient_attention_test.py",
   ]
+if not _has_mlx:
+  collect_ignore.append("tabfm/src/classifier_and_regressor_mlx_test.py")
 # pytorch/model_test.py is a torch<->jax parity test: it imports both flax and
 # torch, so it needs *both* backends installed.
 if not (_has_torch and _has_jax):
   collect_ignore.append("tabfm/src/pytorch/model_test.py")
+# mlx/model_test.py is a torch<->mlx parity test: it needs both backends.
+if not (_has_torch and _has_mlx):
+  collect_ignore.append("tabfm/src/mlx/model_test.py")
 
 
 def pytest_configure(config):  # noqa: D401  (pytest hook)

examples/classification_example.py

@@ -28,6 +28,9 @@ def run_example(model=None) -> np.ndarray:
     # Option B: PyTorch Backend
     # model = tabfm.tabfm_v1_0_0_pytorch.load(model_type="classification")
 
+    # Option C: MLX Backend (Apple silicon)
+    # model = tabfm.tabfm_v1_0_0_mlx.load(model_type="classification")
+
   # 2. Initialize scikit-learn compatible classifier
   clf = tabfm.TabFMClassifier(model=model)
 

examples/regression_example.py

@@ -28,6 +28,9 @@ def run_example(model=None) -> np.ndarray:
     # Option B: PyTorch Backend
     # model = tabfm.tabfm_v1_0_0_pytorch.load(model_type="regression")
 
+    # Option C: MLX Backend (Apple silicon)
+    # model = tabfm.tabfm_v1_0_0_mlx.load(model_type="regression")
+
   # 2. Initialize scikit-learn compatible regressor
   reg = tabfm.TabFMRegressor(model=model)
 

examples/tabarena_classification_example.py

@@ -101,6 +101,9 @@ def run_example(model=None):
     # Option B: PyTorch Backend
     # model = tabfm.tabfm_v1_0_0_pytorch.load(model_type="classification")
 
+    # Option C: MLX Backend (Apple silicon)
+    # model = tabfm.tabfm_v1_0_0_mlx.load(model_type="classification")
+
   x_train, y_train, x_test, y_test = _load_fold_0(TASK_ID)
 
   results = {}

examples/tabarena_regression_example.py

@@ -90,6 +90,9 @@ def run_example(model=None):
     # Option B: PyTorch Backend
     # model = tabfm.tabfm_v1_0_0_pytorch.load(model_type="regression")
 
+    # Option C: MLX Backend (Apple silicon)
+    # model = tabfm.tabfm_v1_0_0_mlx.load(model_type="regression")
+
   x_train, y_train, x_test, y_test = _load_fold_0(TASK_ID)
 
   results = {}

pyproject.toml

@@ -54,8 +54,18 @@ jax = [
     "orbax-checkpoint",
 ]
 pytorch = [
+    # safetensors is required by PyTorchModelHubMixin to load the v1.0.0
+    # weight release; huggingface-hub does not depend on it itself, so a
+    # bare `torch` install fails at load() with a NameError.
+    "safetensors",
     "torch",
 ]
+# MLX backend. Runs on Apple silicon's unified memory (Metal); recent MLX
+# releases also ship CPU/CUDA wheels for Linux. The floor matches the APIs
+# the port relies on (boolean SDPA masks, nan_to_num, Module.set_dtype).
+mlx = [
+    "mlx>=0.31",
+]
 
 # Development deps (unittest, linting, formating,...)
 # Installed through `pip install -e .[dev]`

requirements.txt

@@ -151,6 +151,8 @@ rich==13.7.1
     # via
     #   flax
     #   typer
+safetensors==0.8.0
+    # via tabfm (pyproject.toml)
 scikit-learn==1.6.0
     # via tabfm (pyproject.toml)
 scipy==1.17.1

tabfm/__init__.py

@@ -27,6 +27,12 @@
   # PyTorch is not installed or incomplete, tabfm_v1_0_0_pytorch is not available.
   pass
 
+try:
+  from tabfm.src.mlx import tabfm_v1_0_0 as tabfm_v1_0_0_mlx
+except ImportError:
+  # MLX is not installed or incomplete, tabfm_v1_0_0_mlx is not available.
+  pass
+
 from tabfm.src.classifier_and_regressor import TabFMClassifier, TabFMRegressor
 
 # A new PyPI release will be pushed every time `__version__` is increased.

tabfm/src/classifier_and_regressor.py

@@ -52,6 +52,13 @@
   HAS_TORCH = True
 except ImportError:
   HAS_TORCH = False
+
+try:
+  import mlx.core as mx
+  import mlx.nn as mlx_nn
+  HAS_MLX = True
+except ImportError:
+  HAS_MLX = False
 import pandas as pd
 import scipy.optimize as opt
 import scipy.special
@@ -1826,53 +1833,42 @@ def _predict_step_pytorch(
   return out_t.float().cpu().numpy()  # upcast: numpy has no bfloat16
 
 
-# Compiled predict step functions memoized on the estimators by
-# _batch_forward. They close over nnx.jit state and cannot be pickled.
-_COMPILED_PREDICT_CACHE_ATTRS = (
-    "_predict_step_compiled_with_cat",
-    "_predict_step_compiled_no_cat",
-)
-def _check_classifier_output_dim(output_dim: int, n_classes: int) -> None:
-  """Validates that the model produces logits for all target classes.
-
-  Args:
-    output_dim: Size of the model output's last axis.
-    n_classes: Number of classes seen during fit.
-
-  Raises:
-    ValueError: If the model outputs fewer values per row than there are
-      classes, which usually means a regression checkpoint is being used for
-      classification.
-  """
-  if output_dim < n_classes:
-    raise ValueError(
-        f"The model returned {output_dim} value(s) per row, but this"
-        f" TabFMClassifier was fit on {n_classes} classes. This usually means"
-        " a regression checkpoint is being used for classification. Load the"
-        " classification weights instead, e.g."
-        " `tabfm_v1_0_0.load(model_type='classification')`."
-    )
+def _predict_step_mlx(
+    model: Any,
+    X_batch: np.ndarray,
+    y_batch: np.ndarray,
+    train_size_val: int,
+    ds_batch_val: Optional[np.ndarray],
+    cat_mask_batch: Optional[np.ndarray],
+) -> np.ndarray:
+  """Runs MLX forward pass and returns numpy array."""
+  if not HAS_MLX:
+    raise ImportError("MLX is required to run an MLX model.")
+
+  # MLX arrays live in unified memory; no device transfer is needed.
+  X_m = mx.array(X_batch.astype(np.float32, copy=False))
+  y_m = mx.array(
+      y_batch.astype(np.float32, copy=False)
+      if y_batch.dtype == np.float64
+      else y_batch
+  )
 
+  batch_size = X_batch.shape[0]
+  train_size_m = mx.full((batch_size,), train_size_val, dtype=mx.int32)
 
-def _check_regressor_output_dim(output_dim: int) -> None:
-  """Validates that the model produces scalar regression outputs.
+  if ds_batch_val is not None:
+    d_m = mx.array(ds_batch_val.astype(np.int32, copy=False))
+  else:
+    d_m = mx.full((batch_size,), X_batch.shape[-1], dtype=mx.int32)
 
-  Args:
-    output_dim: Size of the model output's last axis.
+  cat_mask_m = mx.array(cat_mask_batch) if cat_mask_batch is not None else None
 
-  Raises:
-    ValueError: If the model outputs more than one value per row, which
-      usually means a classification checkpoint is being used for regression.
-  """
-  if output_dim != 1:
-    raise ValueError(
-        f"The model returned {output_dim} values per row, but TabFMRegressor"
-        " expects a single regression output. This usually means a"
-        " classification checkpoint is being used for regression; note that"
-        " `tabfm_v1_0_0.load()` defaults to the classification weights. Load"
-        " the regression weights instead with"
-        " `tabfm_v1_0_0.load(model_type='regression')`."
-    )
+  # MLX is lazy and builds no autograd graph unless gradients are requested;
+  # mx.eval materializes the result before the numpy conversion.
+  out_m = model(X_m, y_m, train_size_m, cat_mask=cat_mask_m, d=d_m)
+  out_m = out_m.astype(mx.float32)  # upcast: numpy has no bfloat16
+  mx.eval(out_m)
+  return np.array(out_m)
 
 
 # ---------------------------------------------------------------------------
@@ -2240,9 +2236,11 @@ def _batch_forward(
         where test_size = n_samples - train_size.
     """
     is_torch = HAS_TORCH and isinstance(self.model, torch.nn.Module)
+    is_mlx = HAS_MLX and isinstance(self.model, mlx_nn.Module)
 
-    if is_torch:
-      # --- PyTorch execution path ---
+    if is_torch or is_mlx:
+      # --- Eager (PyTorch / MLX) execution path ---
+      predict_step = _predict_step_pytorch if is_torch else _predict_step_mlx
       batch_size_per_process = self.batch_size or Xs.shape[0]
       n_batches = math.ceil(Xs.shape[0] / batch_size_per_process)
       if n_batches > 1:
@@ -2279,7 +2277,7 @@ def _batch_forward(
               constant_values=-100.0,
           )
 
-        out = _predict_step_pytorch(
+        out = predict_step(
             self.model,
             X_batch,
             y_batch,
@@ -3031,9 +3029,11 @@ def _batch_forward(
       Model outputs of shape (n_datasets, n_test, output_dim).
     """
     is_torch = HAS_TORCH and isinstance(self.model, torch.nn.Module)
+    is_mlx = HAS_MLX and isinstance(self.model, mlx_nn.Module)
 
-    if is_torch:
-      # --- PyTorch execution path ---
+    if is_torch or is_mlx:
+      # --- Eager (PyTorch / MLX) execution path ---
+      predict_step = _predict_step_pytorch if is_torch else _predict_step_mlx
       batch_size_per_process = getattr(self, "batch_size", 1) or Xs.shape[0]
       n_batches = math.ceil(Xs.shape[0] / batch_size_per_process)
       if n_batches > 1:
@@ -3070,7 +3070,7 @@ def _batch_forward(
               constant_values=-100.0,
           )
 
-        out = _predict_step_pytorch(
+        out = predict_step(
             self.model,
             X_batch,
             y_batch,