Add MLX backend for Apple silicon

[ovuruska]

Summary

Adds an MLX backend (tabfm/src/mlx/, pip install -e .[mlx]) for native Apple-silicon inference, following the existing backend pattern (model.py + tabfm_v1_0_0.py).

The port mirrors the PyTorch module/parameter naming and Linear [out, in] layout, so it loads the PyTorch v1.0.0 safetensors release (google/tabfm-1.0.0-pytorch) directly — no weight conversion step and no separate checkpoint release needed.

Also fixes a pre-existing bug found while testing: the pytorch extra was missing safetensors, so tabfm_v1_0_0_pytorch.load() failed with a NameError on a bare .[pytorch] install.

What's included

Integration point	Change
tabfm/src/mlx/model.py	Faithful forward-path port (RMSNorm / Fourier / PerDimScale / RoPE phases in fp32, SDPA at scale=1.0, checkpoint-loaded RoPE freqs)
tabfm/src/mlx/tabfm_v1_0_0.py	load() with HF download, process-wide cache, bf16 default
tabfm/__init__.py	Guarded tabfm_v1_0_0_mlx export
classifier_and_regressor.py	HAS_MLX flag, _predict_step_mlx, MLX dispatch in both _batch_forwards (the PyTorch branch is generalized to a shared eager branch; JAX path untouched)
pyproject.toml	mlx extra; safetensors added to the pytorch extra (bug fix)
conftest.py	Skips MLX tests when mlx (or torch, for the parity test) is not installed
CI	Installs the mlx extra (manylinux cp311 wheels exist) so the new tests run on ubuntu-latest
Docs	README install/quick-start Option C, example scripts, CHANGELOG [Unreleased]

Numerical fidelity

• torch<->mlx parity test (tabfm/src/mlx/model_test.py): max abs diff < 1e-4 in float32 for both classification and regression, with zero-initialized tensors randomized so all paths (Fourier buffers, PerDimScale, cls/ind tokens) are exercised.
• The tanh-approximation gelu (matching jax.nn.gelu) is implemented manually since MLX's built-in gelu variants use erf/sigmoid approximations; MLX has no one_hot, so it is built by comparison against the class range.

Released-checkpoint compatibility

• Safetensors-header level: all tensor names and shapes of the classification (913 tensors) and regression (918 tensors) checkpoints map one-to-one onto the MLX module tree (strict load_weights).
• End-to-end with the real 6.5GB classification checkpoint (Apple M1 Pro, 16 GB): MLX and PyTorch backends produce identical predictions on a mixed-type dataset; max class-probability difference 4.3e-3 under bf16 compute.

Performance (measured)

Real v1.0.0 classification weights; 200 train / 50 test rows, 8 numeric + 2 categorical features, n_estimators=8; each backend in a fresh process. "Warm" = second predict_proba after fit.

Backend	Model load	Cold fit+predict	Warm predict	Accuracy
MLX (Metal)	3.3 s (mmap + eager materialize)	9.5 s	3.1 s	0.940
PyTorch (MPS)	52.5 s	14.1 s	5.9 s	0.920
PyTorch (CPU)	22.9 s	48.0 s	53.3 s	0.920

Warm inference: ~1.9× faster than torch-MPS, ~17× faster than torch-CPU. Time-to-first-prediction: 12.8 s vs 66.6 s (torch-MPS incl. load) — safetensors mmap makes the load step cheap (weights are now materialized eagerly at load() per review feedback). The 0.94-vs-0.92 accuracy delta is one borderline sample out of 50 (bf16 noise), consistent with the 4.3e-3 probability agreement.

vs. classic tabular ML baselines (same dataset, default hyperparameters)

Model	Fit	Predict	Accuracy
TabFM v1.0.0 (MLX, zero-shot)	no gradient training (in-context)	3.2 s (warm)	0.940
XGBoost 3.2	0.10 s	0.003 s	0.940
sklearn RandomForest	0.05 s	0.003 s	0.900
sklearn HistGradientBoosting	0.54 s	0.016 s	0.860

Does tuning close the gap? On a harder task (nonlinear target with a categorical interaction, same 200/50 split, 5 dataset seeds), zero-shot TabFM beats even a tuned XGBoost on every seed:

Model (5-seed mean)	Accuracy	Wall clock / run
TabFM v1.0.0 (MLX, zero-shot)	0.944 ± 0.032	14.7 s
XGBoost tuned (RandomizedSearchCV, 30 iter × 3-fold)	0.856 ± 0.060	2.2 s
XGBoost default	0.852 ± 0.043	0.2 s

Hyperparameter search moves XGBoost by only +0.004 at 200 training rows — the small-data regime is exactly where in-context learning pays off. On the easier linearly-separable dataset above, both tie at 0.940.

Honest framing: gradient-boosted trees remain far faster on small tabular data. TabFM matches XGBoost's top accuracy here zero-shot — no per-dataset training or tuning (fit() only prepares encoders/ensemble views). This table is a speed/behavior sanity check; dataset-quality comparisons are the TabArena evals in results/*.parquet.

All numbers benchmarked on an Apple M1 Pro (16 GB).

Test plan

• pytest -vv tabfm/src/mlx/model_test.py tabfm/src/classifier_and_regressor_mlx_test.py — 4 passed (Apple M1 Pro, macOS, Python 3.11, mlx 0.31.2, torch 2.12.1)
• Full suite with jax+pytorch+mlx installed: 76 passed, 0 failed
• Safetensors header vs. MLX module tree: 0 missing / 0 extra / 0 shape mismatches for both tasks
• End-to-end load() + fit/predict/predict_proba with the real classification checkpoint, MLX vs PyTorch cross-check (max prob diff 4.3e-3, identical predictions)

🤖 Generated with Claude Code

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.