diff --git a/.github/workflows/pytest_and_autopublish.yml b/.github/workflows/pytest_and_autopublish.yml index fc1a3f0..1caccd9 100644 --- a/.github/workflows/pytest_and_autopublish.yml +++ b/.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) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1782ea4..8fab521 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/README.md b/README.md index 0fcffd5..cf877fa 100644 --- a/README.md +++ b/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: diff --git a/conftest.py b/conftest.py index a651128..04997aa 100644 --- a/conftest.py +++ b/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) diff --git a/examples/classification_example.py b/examples/classification_example.py index fc9522b..79002a5 100644 --- a/examples/classification_example.py +++ b/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) diff --git a/examples/regression_example.py b/examples/regression_example.py index 51f4e2f..809fbad 100644 --- a/examples/regression_example.py +++ b/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) diff --git a/examples/tabarena_classification_example.py b/examples/tabarena_classification_example.py index 1d38821..903c7bc 100644 --- a/examples/tabarena_classification_example.py +++ b/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 = {} diff --git a/examples/tabarena_regression_example.py b/examples/tabarena_regression_example.py index 01c3864..f762cdc 100644 --- a/examples/tabarena_regression_example.py +++ b/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 = {} diff --git a/pyproject.toml b/pyproject.toml index 784274d..2938dcb 100644 --- a/pyproject.toml +++ b/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]` diff --git a/requirements.txt b/requirements.txt index 9a5e4e4..6068f94 100644 --- a/requirements.txt +++ b/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 diff --git a/tabfm/__init__.py b/tabfm/__init__.py index cb71100..2c15d98 100644 --- a/tabfm/__init__.py +++ b/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. diff --git a/tabfm/src/classifier_and_regressor.py b/tabfm/src/classifier_and_regressor.py index 51dc867..9f0037c 100644 --- a/tabfm/src/classifier_and_regressor.py +++ b/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, diff --git a/tabfm/src/classifier_and_regressor_mlx_test.py b/tabfm/src/classifier_and_regressor_mlx_test.py new file mode 100644 index 0000000..432e83a --- /dev/null +++ b/tabfm/src/classifier_and_regressor_mlx_test.py @@ -0,0 +1,102 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest +import numpy as np + +from tabfm.src.mlx import model as mlx_model +from tabfm.src.classifier_and_regressor import TabFMClassifier, TabFMRegressor + + +class MLXClassifierRegressorTest(unittest.TestCase): + + def test_classifier_fit_predict(self): + np.random.seed(42) + # Instantiate small MLX model + model = mlx_model.TabFM( + embed_dim=8, + max_classes=3, + 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=True, + ) + + clf = TabFMClassifier( + model=model, + n_estimators=2, + batch_size=2, + random_state=42, + ) + + X = np.random.rand(10, 3) + y = np.random.randint(0, 3, size=10) + + clf.fit(X, y) + + # Predict + preds = clf.predict(X) + self.assertEqual(preds.shape, (10,)) + self.assertTrue(np.all(preds >= 0) and np.all(preds < 3)) + + # Predict proba + probs = clf.predict_proba(X) + self.assertEqual(probs.shape, (10, 3)) + np.testing.assert_allclose(np.sum(probs, axis=1), 1.0, rtol=1e-5) + + def test_regressor_fit_predict(self): + np.random.seed(42) + # Instantiate small MLX model + model = mlx_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, + ) + + reg = TabFMRegressor( + model=model, + n_estimators=2, + batch_size=2, + random_state=42, + ) + + X = np.random.rand(10, 3) + y = np.random.rand(10) + + reg.fit(X, y) + + # Predict + preds = reg.predict(X) + self.assertEqual(preds.shape, (10,)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tabfm/src/mlx/__init__.py b/tabfm/src/mlx/__init__.py new file mode 100644 index 0000000..a81104e --- /dev/null +++ b/tabfm/src/mlx/__init__.py @@ -0,0 +1,15 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# tabfm src mlx package diff --git a/tabfm/src/mlx/model.py b/tabfm/src/mlx/model.py new file mode 100644 index 0000000..be25778 --- /dev/null +++ b/tabfm/src/mlx/model.py @@ -0,0 +1,550 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Faithful MLX port of the TabFM architecture (forward path only). + +Module/param names mirror the PyTorch port (tabfm/src/pytorch/model.py), which +in turn mirrors the JAX model, so the pre-trained PyTorch safetensors +checkpoint loads mechanically: MLX flattens list-of-module children with the +same ``blocks.0.attn.q_proj.weight`` key scheme as ``torch.nn.Module`` and +``mlx.nn.Linear`` stores its weight as ``[out, in]`` like PyTorch, so no +transposes or key remapping are required. Checkpoint buffers (``rope.freqs``, +``fourier_frequencies``) are plain ``mx.array`` attributes, which MLX treats +as loadable parameters. + +Numerical-fidelity notes (see the parity test in model_test.py): +- RMSNorm, the Fourier expansion, PerDimScale's softplus and RoPE's phase + computation all run in float32 with a cast back to the compute dtype, + matching the JAX/PyTorch implementations. +- Attention uses SDPA at scale=1.0 (the 1/sqrt(d) factor is folded into + PerDimScale) with q/k RMSNorm applied after RoPE. +- RoPE inverse frequencies are LOADED from the checkpoint, not recomputed. +""" + +import math +from typing import List, Optional + +import mlx.core as mx +import mlx.nn as nn + + +def _gelu_tanh(x): + # jax.nn.gelu defaults to the tanh approximation -> match it exactly + # (mlx.nn's gelu variants use erf / sigmoid approximations instead). + return 0.5 * x * ( + 1.0 + mx.tanh(0.7978845608028654 * (x + 0.044715 * x * x * x)) + ) + + +def get_activation(name): + return {"relu": nn.relu, "gelu": _gelu_tanh, "silu": nn.silu}[name] + + +def _softplus_f32(x): + # logaddexp(x, 0) == log(1 + exp(x)) computed stably in float32. + return mx.logaddexp(x.astype(mx.float32), mx.array(0.0, dtype=mx.float32)) + + +class RMSNorm(nn.Module): + + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.weight = mx.ones((dim,)) + self.eps = eps + + def __call__(self, x): + # Normalize entirely in float32 (x * rsqrt * weight), cast back at the + # end -- matches JAX/Flax and the PyTorch port. Doing the multiply in bf16 + # loses precision and accumulates across the ~36 RMSNorms per stack. + dt = x.dtype + xf = x.astype(mx.float32) + v = mx.mean(xf * xf, axis=-1, keepdims=True) + out = (xf * mx.rsqrt(v + self.eps)) * self.weight.astype(mx.float32) + return out.astype(dt) + + +def _rotate_with_freqs(x, freqs): + """Interleaved RoPE over the T axis of [B, T, N, Dh] (lucidrains convention).""" + t = x.shape[1] + # Phases in float32; cos/sin cast to the compute dtype afterwards (matches + # the PyTorch port, which upcasts the checkpoint-loaded freqs the same way). + f = ( + mx.arange(t).astype(mx.float32)[:, None] + * freqs.astype(mx.float32)[None, :] + ) + cos = mx.repeat(mx.cos(f), 2, axis=-1)[None, :, None, :].astype(x.dtype) + sin = mx.repeat(mx.sin(f), 2, axis=-1)[None, :, None, :].astype(x.dtype) + x1, x2 = x[..., 0::2], x[..., 1::2] + rot = mx.stack([-x2, x1], axis=-1).reshape(x.shape) + return x * cos + rot * sin + + +def rope_interleaved(x, base): + """RoPE fallback that recomputes the inverse frequencies from ``base``.""" + dh = x.shape[-1] + inv = 1.0 / (base ** (mx.arange(0, dh, 2).astype(mx.float32) / dh)) + return _rotate_with_freqs(x, inv) + + +class RoPE(nn.Module): + """One RoPE per Encoder, holding the inverse-frequency buffer loaded FROM the + checkpoint (JAX stores `rope.freqs`, computed in bf16 at train time -- + recomputing it in fp32 differs by ~1e-3 and that error grows with sequence + length).""" + + def __init__(self, dim, base): + super().__init__() + # init = formula; overwritten on load + self.freqs = 1.0 / (base ** (mx.arange(0, dim, 2).astype(mx.float32) / dim)) + + def rotate(self, x): # x: [B, T, N, Dh], rotate over the T axis + return _rotate_with_freqs(x, self.freqs) + + +class MultiheadAttention(nn.Module): + + def __init__(self, d_model, nhead, rope_base=None): + super().__init__() + self.nhead, self.hd = nhead, d_model // nhead + self.rope_base = rope_base # None => no RoPE + self.q_proj = nn.Linear(d_model, d_model) + self.k_proj = nn.Linear(d_model, d_model) + self.v_proj = nn.Linear(d_model, d_model) + self.out_proj = nn.Linear(d_model, d_model) + self.query_ln = RMSNorm(self.hd) + self.key_ln = RMSNorm(self.hd) + self.per_dim_scale = mx.zeros((self.hd,)) + + def __call__(self, query, key, value, attn_mask=None, rope=None): + b, tq, d = query.shape + q = self.q_proj(query).reshape(b, tq, self.nhead, self.hd) + k = self.k_proj(key).reshape(b, key.shape[1], self.nhead, self.hd) + v = self.v_proj(value).reshape(b, value.shape[1], self.nhead, self.hd) + if self.rope_base is not None: + # Use the Encoder's shared RoPE (checkpoint-loaded freqs) when provided; + # fall back to recomputing only if absent. + if rope is not None: + q, k = rope.rotate(q), rope.rotate(k) + else: + q = rope_interleaved(q, self.rope_base) + k = rope_interleaved(k, self.rope_base) + q, k = self.query_ln(q), self.key_ln(k) + # per-dim scale in float32 (softplus), then cast to compute dtype -- + # matches JAX PerDimScale. + scale = 1.442695041 / math.sqrt(self.hd) * _softplus_f32(self.per_dim_scale) + q = q * scale.astype(q.dtype) + q, k, v = (z.transpose(0, 2, 1, 3) for z in (q, k, v)) # [B,N,T,D] + # Boolean masks (True = attend) are supported natively; SDPA does its + # softmax in float32 internally. + o = mx.fast.scaled_dot_product_attention( + q, k, v, scale=1.0, mask=attn_mask + ) + return self.out_proj(o.transpose(0, 2, 1, 3).reshape(b, tq, d)) + + +class MultiheadAttentionBlock(nn.Module): + + def __init__(self, d_model, nhead, dim_ff, activation="swiglu", + rope_base=None): + super().__init__() + self.attn = MultiheadAttention(d_model, nhead, rope_base) + self.pre_attn_ln = RMSNorm(d_model) + self.post_attn_ln = RMSNorm(d_model) + self.pre_ff_ln = RMSNorm(d_model) + self.post_ff_ln = RMSNorm(d_model) + self.swiglu = activation == "swiglu" + self.linear1 = nn.Linear(d_model, dim_ff) + if self.swiglu: + self.linear1_gate = nn.Linear(d_model, dim_ff) + self.act = nn.silu + else: + self.act = get_activation(activation) + self.linear2 = nn.Linear(dim_ff, d_model) + self.ffn_chunk_size = None # set to an int to chunk the FFN over tokens + + def _ff_impl(self, x): + xn = self.pre_ff_ln(x) + if self.swiglu: + x = self.act(self.linear1_gate(xn)) * self.linear1(xn) + else: + x = self.act(self.linear1(xn)) + return self.post_ff_ln(self.linear2(x)) + + def _ff(self, x): + # FFN chunking: process tokens in slices so the expanded + # [tokens, dim_feedforward] activation is never materialized in full. + if self.ffn_chunk_size is None: + return self._ff_impl(x) + shape = x.shape + flat = x.reshape(-1, shape[-1]) + parts = [ + self._ff_impl(flat[s : s + self.ffn_chunk_size]) + for s in range(0, flat.shape[0], self.ffn_chunk_size) + ] + return mx.concatenate(parts, axis=0).reshape(shape) + + def __call__(self, q, k=None, v=None, attn_mask=None, rope=None): + k = q if k is None else k + v = q if v is None else v + a = self.post_attn_ln( + self.attn( + self.pre_attn_ln(q), + self.pre_attn_ln(k), + self.pre_attn_ln(v), + attn_mask, + rope=rope, + ) + ) + x = q + a + return x + self._ff(x) + + +class InducedSelfAttentionBlock(nn.Module): + + def __init__(self, d_model, nhead, dim_ff, num_inds, activation="swiglu"): + super().__init__() + self.ind_vectors = mx.zeros((num_inds, d_model)) + self.mab1 = MultiheadAttentionBlock(d_model, nhead, dim_ff, activation) + self.mab2 = MultiheadAttentionBlock(d_model, nhead, dim_ff, activation) + + def __call__(self, src, attn_mask=None): + ind = mx.broadcast_to( + self.ind_vectors[None], (src.shape[0],) + self.ind_vectors.shape + ) + hidden = self.mab1(ind, src, src, attn_mask=attn_mask) + return self.mab2(src, hidden, hidden) + + +class Encoder(nn.Module): + + def __init__(self, num_blocks, d_model, nhead, dim_ff, activation="swiglu", + rope_base=100000.0): + super().__init__() + # One RoPE per Encoder (mirrors JAX `tf_row.rope.freqs`), shared by all + # blocks. + self.rope = ( + RoPE(d_model // nhead, rope_base) if rope_base is not None else None + ) + self.blocks = [ + MultiheadAttentionBlock(d_model, nhead, dim_ff, activation, rope_base) + for _ in range(num_blocks) + ] + + def __call__(self, x, attn_mask=None): + for blk in self.blocks: + x = blk(x, attn_mask=attn_mask, rope=self.rope) + return x + + +class SetTransformer(nn.Module): + + def __init__(self, num_blocks, d_model, nhead, dim_ff, num_inds, + activation="swiglu"): + super().__init__() + self.blocks = [ + InducedSelfAttentionBlock(d_model, nhead, dim_ff, num_inds, activation) + for _ in range(num_blocks) + ] + + def __call__(self, src, attn_mask=None): + for blk in self.blocks: + src = blk(src, attn_mask=attn_mask) + return src + + +class MLP(nn.Module): + + def __init__(self, in_dim, hidden_dims: List[int], out_dim, + activation="gelu"): + super().__init__() + self.act = get_activation(activation) + dims = [in_dim] + list(hidden_dims) + self.layers = [] # only Linears; activation applied between + for i in range(len(hidden_dims)): + self.layers.append(nn.Linear(dims[i], dims[i + 1])) + self.layers.append(nn.Linear(dims[-1], out_dim)) + + def __call__(self, x): + for i, lin in enumerate(self.layers): + x = lin(x) + if i < len(self.layers) - 1: + x = self.act(x) + return x + + +class OneHotAndLinear(nn.Module): + + def __init__(self, num_classes, embed_dim): + super().__init__() + self.num_classes = num_classes + self.projection = nn.Linear(num_classes, embed_dim) + + def __call__(self, y): # y: [B, T] int + y_int = y.astype(mx.int32) + # No mx.one_hot; compare against the class range instead. Out-of-range + # labels (e.g. the -100 padding sentinel) match no column and contribute + # only the projection bias -- identical to the PyTorch remap-then-slice. + oh = (y_int[..., None] == mx.arange(self.num_classes)).astype( + self.projection.weight.dtype + ) + return self.projection(oh) + + +class CellEmbedder(nn.Module): + + def __init__(self, embed_dim, max_classes, feature_group_size=3, num_freq=32, + is_classifier=True): + super().__init__() + self.embed_dim = embed_dim + self.fgs = feature_group_size + self.is_classifier = is_classifier + in_dim = feature_group_size + self.fourier_frequencies = mx.zeros((in_dim, num_freq)) + self.fourier_frequencies_cat = mx.zeros((in_dim, num_freq)) + self.in_linear = nn.Linear(num_freq * 2, embed_dim) + self.in_linear_cat = nn.Linear(num_freq * 2, embed_dim) + if is_classifier: # classification: embedding lookup over class ids + self.y_embedder_lookup = nn.Embedding(max_classes, embed_dim) + else: # regression: MLP over the scalar target (y_col_embedder_encoder_nhid=6) + self.y_embedder_lookup = MLP(1, [6], embed_dim, activation="gelu") + self.row_chunk_size = None # chunk the Fourier expansion over rows + + def _group(self, x, d=None): # x: [B,T,H] -> [B,T,H,G] + h = x.shape[-1] + idxs = mx.arange(h) + stacked = [] + if d is not None: + # Per-batch wrap-around over each member's ACTIVE feature count d (not + # the padded width h). Mirrors the JAX `% d_safe` path so zero-padded + # slots are filled with wrapped real features rather than mixing padding + # into groups. + d_safe = mx.maximum(d.astype(mx.int32), 1) # [B] + for i in range(self.fgs): + offset = (2 ** i) - 1 + idx = (idxs[None, :] + offset) % d_safe[:, None] # [B, H] + idx = mx.broadcast_to( + idx[:, None, :], (x.shape[0], x.shape[1], h) + ) # [B, T, H] + stacked.append(mx.take_along_axis(x, idx, axis=-1)) + else: + for i in range(self.fgs): + offset = (2 ** i) - 1 + stacked.append(mx.take(x, (idxs + offset) % h, axis=-1)) + return mx.stack(stacked, axis=-1) + + def _cell(self, x, cat_mask, d=None): + # [B,t,H] -> [B,t,HC,E] (Fourier expansion + sum over G). + # float32 Fourier: args g*freq reach ~30, so sin/cos must run in fp32 + # (matches JAX, whose freq params stay float32). Cast the fourier features + # back to compute dtype before in_linear. + g = self._group(x, d=d)[..., None].astype(mx.float32) + dt = x.dtype + ff = self.fourier_frequencies.astype(mx.float32) + ffc = self.fourier_frequencies_cat.astype(mx.float32) + num_out = self.in_linear( + mx.concatenate([mx.sin(g * ff), mx.cos(g * ff)], axis=-1).astype(dt) + ) + if cat_mask is not None: + cat_out = self.in_linear_cat( + mx.concatenate([mx.sin(g * ffc), mx.cos(g * ffc)], axis=-1).astype(dt) + ) + cmg = self._group( + cat_mask[:, None, :].astype(mx.float32), d=d + ).astype(mx.bool_)[..., None] + return mx.where(cmg, cat_out, num_out).sum(axis=-2) + return num_out.sum(axis=-2) + + def __call__(self, x, y, train_size, cat_mask=None, d=None): + # The Fourier expansion materializes [B,T,HC,G,E]; chunk over rows so that + # huge intermediate never exists in full (rows are independent here). + if self.row_chunk_size is None: + cell = self._cell(x, cat_mask, d=d) + else: + parts = [ + self._cell(x[:, s : s + self.row_chunk_size], cat_mask, d=d) + for s in range(0, x.shape[1], self.row_chunk_size) + ] + cell = mx.concatenate(parts, axis=1) + if self.is_classifier: + num_embeddings = self.y_embedder_lookup.weight.shape[0] + y_clean = mx.clip(y.astype(mx.int32), 0, num_embeddings - 1) + y_emb = self.y_embedder_lookup(y_clean) # [B,T,E] + else: + y_emb = self.y_embedder_lookup(y[..., None].astype(cell.dtype)) + t = x.shape[1] + tm = (mx.arange(t)[None, :] < train_size[:, None])[..., None, None] + out = mx.where(tm, cell + y_emb[:, :, None, :], cell) + if d is not None: + # Zero the padded feature columns (cols >= d): the % d wrap above fills + # them with real features for valid indexing, but they must not enter + # attention. + hc = out.shape[2] + colmask = (mx.arange(hc)[None, :] < d[:, None])[:, None, :, None] + out = mx.where(colmask, out, mx.zeros_like(out)) + return out + + +class ColEmbedding(nn.Module): + + def __init__(self, d_model, num_blocks, nhead, dim_ff, num_inds): + super().__init__() + self.tf_col = SetTransformer(num_blocks, d_model, nhead, dim_ff, num_inds) + self.out_w = nn.Linear(d_model, d_model) + self.ln_w = RMSNorm(d_model) + self.col_chunk_size = None # chunk the independent column axis (B*HC) + + def _stage(self, src, mask): + return self.ln_w(self.out_w(self.tf_col(src, attn_mask=mask))) + + def __call__(self, x, train_size): # x: [B,T,HC,E] + b, t, hc, e = x.shape + src = x.transpose(0, 2, 1, 3).reshape(b * hc, t, e) # [B*HC, T, E] + ts = mx.repeat(train_size, hc, axis=0) # [B*HC] + mask = (mx.arange(t)[None, :] < ts[:, None])[:, None, None, :] + cc = self.col_chunk_size + if cc is None or src.shape[0] <= cc: + out = self._stage(src, mask) + else: + out = mx.concatenate( + [ + self._stage(src[s : s + cc], mask[s : s + cc]) + for s in range(0, src.shape[0], cc) + ], + axis=0, + ) + return out.reshape(b, hc, t, e).transpose(0, 2, 1, 3) + + +class RowInteraction(nn.Module): + + def __init__(self, d_model, num_blocks, nhead, dim_ff, num_cls, + rope_base=100000.0, output_full=True): + super().__init__() + self.tf_row = Encoder(num_blocks, d_model, nhead, dim_ff, + rope_base=rope_base) + self.out_ln = RMSNorm(d_model) + self.num_cls = num_cls + self.output_full = output_full + self.row_chunk_size = None # chunk the independent row axis (B*T) + + def _stage(self, src, mask=None): + out = self.tf_row(src, attn_mask=mask) + return self.out_ln(out if self.output_full else out[:, : self.num_cls, :]) + + def __call__(self, x, d=None): # x: [B,T,HC,E] + b, t, hc, e = x.shape + src = x.reshape(b * t, hc, e) + # Mask cross-column attention to the valid columns (CLS + d real + # features); padded columns (>= d + num_cls) must not be attended to. + # Matches JAX. + mask = None + if d is not None: + d_padded = d.astype(mx.int32) + self.num_cls # [B] + valid = mx.arange(hc)[None, :] < d_padded[:, None] # [B, HC] + mask = mx.repeat(valid, t, axis=0)[:, None, None, :] # [B*T, 1, 1, HC] + rc = self.row_chunk_size + if rc is None or src.shape[0] <= rc: + out = self._stage(src, mask) + else: + out = mx.concatenate( + [ + self._stage( + src[s : s + rc], None if mask is None else mask[s : s + rc] + ) + for s in range(0, src.shape[0], rc) + ], + axis=0, + ) + if self.output_full: + return out.reshape(b, t, hc, e) + return out.reshape(b, t, -1) + + +class ICLearning(nn.Module): + + def __init__(self, d_model, num_blocks, nhead, max_classes, dim_ff, + decoder_hidden, is_classifier=True): + super().__init__() + # ICL has no RoPE. + self.tf_icl = Encoder(num_blocks, d_model, nhead, dim_ff, rope_base=None) + self.ln = RMSNorm(d_model) + self.is_classifier = is_classifier + if is_classifier: # one-hot y-encode; decode to per-class logits + self.y_encoder = OneHotAndLinear(max_classes, d_model) + self.decoder = MLP(d_model, [decoder_hidden], max_classes) + else: # MLP y-encode the scalar target; decode to a single value + self.y_encoder = MLP(1, [decoder_hidden], d_model) + self.decoder = MLP(d_model, [decoder_hidden], 1) + + def __call__(self, reps, y, train_size): # reps: [B,T,d_model] + b, t, _ = reps.shape + tm = mx.arange(t)[None, :] < train_size[:, None] + if self.is_classifier: + y_enc = self.y_encoder(y) + else: + y_enc = self.y_encoder(y[..., None].astype(reps.dtype)) + r = reps + y_enc * tm[..., None] + mask = tm[:, None, None, :] + out = self.tf_icl(r, attn_mask=mask) + return self.decoder(self.ln(out)) + + +class TabFM(nn.Module): + + def __init__(self, *, embed_dim=8, max_classes=3, col_num_blocks=2, + col_nhead=2, col_num_inds=4, row_num_blocks=2, row_nhead=2, + row_num_cls=2, icl_num_blocks=2, icl_nhead=2, ff_factor=2, + feature_group_size=3, num_freq=32, decoder_hidden=None, + is_classifier=True): + super().__init__() + self.max_classes = max_classes + self.is_classifier = is_classifier + ff = embed_dim * ff_factor + icl_dim = embed_dim * row_num_cls + self.cell_embedder = CellEmbedder(embed_dim, max_classes, + feature_group_size, num_freq, + is_classifier) + self.col_embedder = ColEmbedding(embed_dim, col_num_blocks, col_nhead, ff, + col_num_inds) + self.col_embedder_2 = ColEmbedding(embed_dim, col_num_blocks, col_nhead, + ff, col_num_inds) + self.row_interactor = RowInteraction(embed_dim, row_num_blocks, row_nhead, + ff, row_num_cls, output_full=True) + self.row_interactor_2 = RowInteraction(embed_dim, row_num_blocks, + row_nhead, ff, row_num_cls, + output_full=False) + self.cls_tokens = mx.zeros((row_num_cls, embed_dim)) + self.icl_predictor = ICLearning(icl_dim, icl_num_blocks, icl_nhead, + max_classes, icl_dim * ff_factor, + decoder_hidden or icl_dim * 2, + is_classifier) + + def __call__(self, x, y, train_size, cat_mask=None, d=None): + # Mirror the JAX model's entry: replace NaN with the -100 sentinel and + # cast to the compute dtype (JAX: + # `jnp.nan_to_num(X, nan=-100.0).astype(self.dtype)`). NaN is already + # imputed in the shared preprocessing, so nan_to_num is a no-op in the + # normal flow, but it keeps the model robust + JAX-faithful. + x = mx.nan_to_num(x, nan=-100.0).astype(self.cls_tokens.dtype) + emb = self.cell_embedder(x, y, train_size, cat_mask, d=d) + emb = self.col_embedder(emb, train_size) + b, t, _, e = emb.shape + cls = mx.broadcast_to( + self.cls_tokens[None, None], (b, t) + self.cls_tokens.shape + ) + emb = mx.concatenate([cls, emb], axis=2) + emb = self.row_interactor(emb, d=d) + emb = self.col_embedder_2(emb, train_size) + reps = self.row_interactor_2(emb, d=d) + return self.icl_predictor(reps, y, train_size) diff --git a/tabfm/src/mlx/model_test.py b/tabfm/src/mlx/model_test.py new file mode 100644 index 0000000..156af8b --- /dev/null +++ b/tabfm/src/mlx/model_test.py @@ -0,0 +1,165 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest +import numpy as np +import mlx.core as mx +import torch + +from tabfm.src.mlx import model as MLXTabFM +from tabfm.src.pytorch import model as PyTorchTabFM + + +def _randomized_torch_state_dict(torch_model, seed=0): + """Returns the torch state dict with zero-initialized tensors randomized. + + cls_tokens, ind_vectors, per_dim_scale and the Fourier-frequency buffers are + zero at init; randomizing them makes the parity test exercise those paths + (RoPE freqs and Linear weights are already non-trivially initialized). + """ + rng = np.random.default_rng(seed) + state_dict = { + **dict(torch_model.named_parameters()), + **dict(torch_model.named_buffers()), + } + out = {} + for key, value in state_dict.items(): + tensor = value.detach().clone() + if torch.count_nonzero(tensor) == 0: + tensor = torch.tensor( + rng.normal(scale=0.5, size=tuple(tensor.shape)).astype(np.float32) + ) + out[key] = tensor + return out + + +class MLXModelTest(unittest.TestCase): + + def test_mlx_model_instantiation(self): + """Verifies that the MLX model instantiates with config and runs a forward pass.""" + model = MLXTabFM.TabFM( + embed_dim=16, + max_classes=10, + col_num_blocks=2, + col_nhead=2, + col_num_inds=8, + row_num_blocks=2, + row_nhead=2, + row_num_cls=4, + icl_num_blocks=2, + icl_nhead=2, + ff_factor=2, + feature_group_size=3, + is_classifier=True, + ) + self.assertIsNotNone(model) + + # Run dummy forward pass + np.random.seed(0) + x = mx.array(np.random.randn(2, 4, 6).astype(np.float32)) # [B, T, H] + y = mx.array(np.random.randint(0, 3, size=(2, 4))) # [B, T] + train_size = mx.array(np.array([2, 3], dtype=np.int32)) # [B] + out = model(x, y, train_size) + self.assertEqual(out.shape, (2, 4, 10)) + + def test_torch_mlx_parity(self): + """Verifies PyTorch vs MLX model outputs are numerically equal up to 1e-4.""" + for is_classifier in [True, False]: + with self.subTest(is_classifier=is_classifier): + # 1. Config definition (shared by both backends) + cfg = dict( + embed_dim=32, + max_classes=4, + col_num_blocks=2, + col_nhead=4, + col_num_inds=16, + row_num_blocks=2, + row_nhead=4, + row_num_cls=4, + icl_num_blocks=3, + icl_nhead=4, + ff_factor=4, + feature_group_size=3, + is_classifier=is_classifier, + ) + + # 2. Instantiate PyTorch model (random init) and randomize its + # zero-initialized tensors so all paths are exercised. + torch_model = PyTorchTabFM.TabFM(**cfg) + state_dict = _randomized_torch_state_dict(torch_model, seed=7) + torch_model.load_state_dict(state_dict, strict=True) + torch_model.eval() + + # 3. Load the same weights into the MLX model. Names and layouts + # mirror the PyTorch port, so the state dict maps one-to-one. + mlx_model = MLXTabFM.TabFM(**cfg) + mlx_model.load_weights( + [(k, mx.array(v.numpy())) for k, v in state_dict.items()], + strict=True, + ) + mlx_model.eval() + + # 4. Prepare random input data + b, t, h = 3, 5, 8 + np.random.seed(123) + x_np = np.random.normal(size=(b, t, h)).astype(np.float32) + + if is_classifier: + y_np = np.random.randint( + 0, cfg["max_classes"], size=(b, t) + ).astype(np.float32) + else: + y_np = np.random.normal(size=(b, t)).astype(np.float32) + + train_size_np = np.array([2, 3, 4], dtype=np.int32) + d_np = np.array([5, 6, 7], dtype=np.int32) # active feature counts + cat_mask_np = np.zeros((b, h), dtype=bool) + cat_mask_np[0, :3] = True + cat_mask_np[1, :4] = True + + # 5. Forward passes + with torch.no_grad(): + torch_out = torch_model( + torch.from_numpy(x_np), + torch.from_numpy(y_np), + torch.from_numpy(train_size_np), + cat_mask=torch.from_numpy(cat_mask_np), + d=torch.from_numpy(d_np), + ).numpy() + + mlx_out = np.array( + mlx_model( + mx.array(x_np), + mx.array(y_np), + mx.array(train_size_np), + cat_mask=mx.array(cat_mask_np), + d=mx.array(d_np), + ) + ) + + # 6. Compare PyTorch vs MLX outputs + diff = np.abs(torch_out - mlx_out) + max_diff = np.max(diff) + mean_diff = np.mean(diff) + + self.assertLess( + max_diff, + 1e-4, + f"Fidelity discrepancy found: max diff = {max_diff}, mean diff =" + f" {mean_diff} for is_classifier={is_classifier}", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tabfm/src/mlx/tabfm_v1_0_0.py b/tabfm/src/mlx/tabfm_v1_0_0.py new file mode 100644 index 0000000..009dfdb --- /dev/null +++ b/tabfm/src/mlx/tabfm_v1_0_0.py @@ -0,0 +1,158 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Loads pre-trained TabFM v1.0.0 weights into the MLX backend. + +The MLX port shares its module/parameter naming with the PyTorch port, so it +loads the PyTorch safetensors checkpoint (google/tabfm-1.0.0-pytorch) directly +via ``mx.load`` -- no separate MLX weight release or key remapping is needed. +""" + +import json +import os +import threading +from typing import Any, Dict, Optional + +from absl import logging +import mlx.core as mx +from huggingface_hub import snapshot_download + +from tabfm.src.mlx.model import TabFM + +# The MLX backend reuses the PyTorch weight release: parameter names and +# layouts (Linear as [out, in]) are identical, so the safetensors file loads +# one-to-one. +HF_REPO_ID = "google/tabfm-1.0.0-pytorch" + +_CONFIG_NAME = "config.json" +_WEIGHTS_NAME = "model.safetensors" + +_LOAD_CACHE_LOCK = threading.Lock() +_LOAD_CACHE: Dict[Any, TabFM] = {} + + +def _model_kwargs_from_config(config: Dict[str, Any]) -> Dict[str, Any]: + """Translates a checkpoint config.json into TabFM constructor kwargs.""" + kwargs = dict(config) + if "task" in kwargs: + kwargs["is_classifier"] = kwargs.pop("task") == "classification" + for key in ("model_type", "version", "framework"): + kwargs.pop(key, None) + return kwargs + + +def _load_from_dir(local_dir: str, model_type: str) -> TabFM: + """Instantiates a TabFM from a directory with config.json + safetensors.""" + weights_path = os.path.join(local_dir, _WEIGHTS_NAME) + if not os.path.exists(weights_path): + raise FileNotFoundError( + f"No {_WEIGHTS_NAME} found in {local_dir}. The MLX backend loads " + "safetensors checkpoints (as released in " + f"https://huggingface.co/{HF_REPO_ID})." + ) + + cfg_path = os.path.join(local_dir, _CONFIG_NAME) + if os.path.exists(cfg_path): + with open(cfg_path) as f: + model_kwargs = _model_kwargs_from_config(json.load(f)) + else: + # no config.json: pass is_classifier explicitly + logging.warning("No config.json found in %s", local_dir) + model_kwargs = {"is_classifier": model_type == "classification"} + + model = TabFM(**model_kwargs) + # strict=True: the checkpoint keys (parameters + buffers such as + # `rope.freqs` and `fourier_frequencies`) must map one-to-one onto the + # MLX module tree. + model.load_weights(list(mx.load(weights_path).items()), strict=True) + return model + + +def load( + model_type: str = "classification", + checkpoint_path: Optional[str] = None, + *, + dtype: Any = mx.bfloat16, + use_cache: bool = True, +) -> TabFM: + """Loads the MLX TabFM v1.0.0 model with pre-trained weights. + + The checkpoint is stored in float32, but the model is designed to run in + bfloat16 (matching the JAX release's ``dtype=jnp.bfloat16`` compute default), + with a few internal fp32 upcasts. ``dtype`` casts the model accordingly; pass + ``None`` to keep the float32 weights. + + ``dtype`` is provided for float32 debugging / quality comparison; the model + is designed for bfloat16 and this option may be removed in a future release. + + Args: + model_type: 'classification' or 'regression'. + checkpoint_path: Local directory with the checkpoint (either the directory + containing the ``model_type`` subfolder, or that subfolder itself). If + None, downloads from Hugging Face (google/tabfm-1.0.0-pytorch). + dtype: Compute dtype to cast the model to after loading. Defaults to + bfloat16; pass None to keep the float32 weights. MLX arrays live in + unified memory, so there is no device argument. + use_cache: Reuse a process-wide cached model for identical settings. + + Returns: + An eval-mode MLX TabFM model with pre-trained weights loaded. + """ + if model_type not in ("classification", "regression"): + raise ValueError( + f"Unsupported model_type: {model_type!r}. " + "Must be 'classification' or 'regression'." + ) + + cache_key = (model_type, checkpoint_path, str(dtype)) + if use_cache: + _LOAD_CACHE_LOCK.acquire() + try: + if use_cache and cache_key in _LOAD_CACHE: + return _LOAD_CACHE[cache_key] + + if checkpoint_path is None: + logging.info( + "Downloading TabFM v1.0.0 %s weights from Hugging Face...", + model_type, + ) + base_path = snapshot_download( + repo_id=HF_REPO_ID, + allow_patterns=[f"{model_type}/**"], + ) + local_dir = os.path.join(base_path, model_type) + else: + local_dir = checkpoint_path + if not os.path.isdir(local_dir): + raise FileNotFoundError(f"Local checkpoint path not found: {local_dir}") + sub = os.path.join(local_dir, model_type) + if os.path.isdir(sub): + local_dir = sub + + model = _load_from_dir(local_dir, model_type) + + if dtype is not None: + model.set_dtype(dtype) # engage the bf16 compute design (see docstring) + # Materialize the (lazy) checkpoint read + dtype cast now, so load errors + # surface here rather than inside the first forward pass, and the cached + # model is safe to share across threads. + mx.eval(model.parameters()) + model.eval() + + if use_cache: + _LOAD_CACHE[cache_key] = model + return model + finally: + if use_cache: + _LOAD_CACHE_LOCK.release()