Add LoRA fine-tuning to the MLX backend

.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,25 @@ To release a new version (e.g. from `1.0.0` -> `2.0.0`):
 
 -->
 
+## [Unreleased]
+
+* Added LoRA fine-tuning utilities for the MLX backend
+  (`tabfm/src/mlx/lora.py`): `apply_lora` / `train_lora` / `fit_lora` /
+  `merge_lora` / adapter save-load. Only the low-rank adapters are trained;
+  the pre-trained base weights stay frozen. `fit_lora` reuses the sklearn
+  wrapper's `fit` preprocessing so adapters train on the exact numeric
+  distribution the model sees at predict time.
+* 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

CLAUDE.md

@@ -0,0 +1,110 @@
+# CLAUDE.md
+
+Guidance for AI coding agents working in this repository.
+
+## Project overview
+
+TabFM is a scikit-learn compatible tabular foundation model that performs
+zero-shot classification and regression via in-context learning. The reference
+implementation is JAX/Flax (nnx); a numerically parity-verified PyTorch port
+exists alongside it. Pre-trained v1.0.0 weights are downloaded from Hugging
+Face Hub at load time.
+
+## Repository layout
+
+```
+tabfm/__init__.py                      # Public API; guarded backend imports
+tabfm/src/classifier_and_regressor.py  # sklearn wrappers + preprocessing + backend dispatch
+tabfm/src/jax/                         # Reference backend (model.py, tabfm_v1_0_0.py, ...)
+tabfm/src/pytorch/                     # PyTorch port (model.py, tabfm_v1_0_0.py)
+tabfm/src/mlx/                         # MLX port (model.py, tabfm_v1_0_0.py)
+tabfm/src/hugging_face/                # Weight conversion / upload utilities
+examples/                              # Runnable end-to-end examples
+conftest.py                            # Skips backend tests when the backend is not installed
+```
+
+## Coding style
+
+Google Python style, enforced by `pyink` (see `[tool.pyink]` in
+`pyproject.toml`):
+
+- **2-space indentation**, 80-column lines, majority double quotes.
+- Every `.py` file starts with the Apache 2.0 license header
+  (`# Copyright 2026 Google LLC ...`) — copy it verbatim from any existing file.
+- Module docstring after the header describing the module's purpose.
+- Google-style docstrings with `Args:` / `Returns:` sections on public
+  functions and methods. Classes document public attributes under
+  `Attributes:`; sklearn fitted attributes use the trailing-underscore
+  convention (`categories_`, `tfm_`) and are also declared as class-level type
+  annotations.
+- Type hints use `typing` (`Optional`, `List`, `Dict`, `Union`, `Any`).
+  Shape-checked signatures in `classifier_and_regressor.py` use `jaxtyping`
+  annotations with the `@jt.typed` decorator.
+- Logging via `absl.logging`, not `print` (except `verbose=True` user-facing
+  output).
+- sklearn-style `X` / `y` capitalization is allowed
+  (`# pylint: disable=invalid-name`).
+- Comments explain **why**, not what — especially numerical-precision
+  rationale (fp32 upcasts, JAX parity). Keep that density when editing model
+  code.
+
+## Backend pattern
+
+Each compute backend lives in `tabfm/src/<backend>/` with the same two files:
+
+- `model.py` — the architecture. Module and parameter names **mirror the JAX
+  model** so weight conversion is mechanical (`cell_embedder`, `col_embedder`,
+  `row_interactor`, `icl_predictor`, `q_proj`, `per_dim_scale`, ...).
+- `tabfm_v1_0_0.py` — a `load(model_type, checkpoint_path, ...)` function that
+  downloads pre-trained weights from Hugging Face Hub, with a process-wide
+  cache keyed on the load arguments.
+
+Integration points when adding a backend:
+
+1. `tabfm/__init__.py`: `try/except ImportError` guarded import exposing
+   `tabfm_v1_0_0_<backend>`.
+2. `classifier_and_regressor.py`: `HAS_<BACKEND>` flag from a guarded import;
+   an `isinstance` check in both `_batch_forward` methods (classifier and
+   regressor); a `_predict_step_<backend>` helper that takes numpy in and
+   returns numpy out.
+3. `pyproject.toml`: an optional-dependency extra named after the backend.
+4. `conftest.py`: add the backend's test files to `collect_ignore` when the
+   backend (or a parity-test dependency) is not installed.
+5. `README.md`: installation + quick-start option for the backend.
+
+## Numerical fidelity rules
+
+The checkpoint is float32; the model is designed to run in **bfloat16** with
+targeted fp32 upcasts. When porting or editing model code, preserve:
+
+- RMSNorm: normalize entirely in float32, cast back at the end.
+- Fourier feature expansion: `sin`/`cos` computed in float32.
+- PerDimScale: softplus in float32, then cast to compute dtype.
+- Attention: SDPA with `scale=1.0` (scaling is folded into PerDimScale);
+  q/k RMSNorm after RoPE.
+- RoPE inverse frequencies are **loaded from the checkpoint**, not recomputed.
+- Model entry: `nan_to_num(x, nan=-100.0)` then cast to compute dtype.
+
+New-backend outputs must match the reference within `1e-4` max abs diff in
+float32 (see the parity tests in `tabfm/src/pytorch/model_test.py`).
+
+## Tests
+
+- `unittest.TestCase` classes in `*_test.py` files colocated with the source,
+  ending with `if __name__ == "__main__": unittest.main()`.
+- Parity tests instantiate small random-init configs of two backends, convert
+  weights, and assert max abs diff < 1e-4.
+- sklearn integration tests (`classifier_and_regressor_<backend>_test.py`)
+  run `fit`/`predict`/`predict_proba` on tiny random data with a small
+  random-init model.
+- Run: `pytest -vv -n auto` from the repo root (CI uses Python 3.11 with
+  `pip install -e .[dev,jax,pytorch,mlx]`).
+- Tests must not require network access or pre-trained weights.
+
+## Releases / docs
+
+- `CHANGELOG.md` follows keepachangelog.com; bumping `__version__` in
+  `tabfm/__init__.py` triggers the PyPI auto-publish workflow — do not bump it
+  in feature PRs.
+- `requirements.txt` is a pip-compile lock for the full reproducible
+  environment; `pyproject.toml` constrains only known incompatibilities.

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
@@ -134,6 +152,28 @@ predictions = reg.predict(X_test)
 print("Predicted Prices:", predictions)
 ```
 
+### 3. LoRA Fine-Tuning (MLX backend, experimental)
+
+The MLX backend supports parameter-efficient fine-tuning: the pre-trained
+weights stay frozen and only low-rank adapters are trained.
+
+```python
+from tabfm import TabFMClassifier, tabfm_v1_0_0_mlx
+from tabfm.src.mlx import lora
+
+model = tabfm_v1_0_0_mlx.load()
+lora.apply_lora(model, rank=8)      # freeze base, add adapters (ICL blocks)
+
+clf = TabFMClassifier(model=model)
+# Fits the usual encoders/ensembles, then trains ONLY the adapters on the
+# same preprocessed matrix the wrapper feeds the model at predict time.
+lora.fit_lora(clf, X_train, y_train, steps=200)
+
+predictions = clf.predict(X_test)   # inference with the adapted model
+
+lora.merge_lora(model)              # optional: fold adapters into the base
+```
+
 ---
 
 ## Examples Directory
@@ -167,6 +207,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,18 @@
       "tabfm/src/jax/checkpointing_test.py",
       "tabfm/src/jax/memory_efficient_attention_test.py",
   ]
+if not _has_mlx:
+  collect_ignore += [
+      "tabfm/src/classifier_and_regressor_mlx_test.py",
+      "tabfm/src/mlx/lora_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.