Add an MLX backend for TabFM v1.0.0 (Apple silicon)

[franz101]

Adds an MLX backend for TabFM v1.0.0, so the model runs on Apple silicon GPUs without PyTorch or JAX. It follows the existing backend structure: tabfm/src/mlx/ mirrors tabfm/src/pytorch/ with identical module and parameter names, so checkpoint conversion is a pure format translation.

Changes

• tabfm/src/mlx/ — model.py (forward path), tabfm_v1_0_0.py (weight loading, bfloat16 with the same fp32 internals as the torch path), convert.py (safetensors -> .npz)
• tabfm/src/classifier_and_regressor.py — _predict_step_mlx / _build_context_cache_mlx mirroring the pytorch equivalents, behind HAS_MLX
• tabfm/__init__.py — guarded tabfm_v1_0_0_mlx import, same try/except ImportError pattern as the torch and JAX backends
• pyproject.toml — optional [mlx] extra (mlx, safetensors)

Every entry point raises ImportError if MLX isn't installed, so nothing changes for existing users.

Tests

21 unit tests covering forward parity, prefill/decode consistency, context lengths that are not multiples of 128, padding edges, int8 KV-cache quantization and decode parity, cache concatenation, cross-member batching, and the load/reconversion path.

On the 128-row padding

prefill/decode in the JAX and PyTorch backends pad the sequence to a multiple of 128 with the -100.0 sentinel and unpad the logits afterwards. That padding is a sharding artifact, and the sentinel rows are excluded everywhere: train_size counts non-sentinel labels, and each stage is per-position, per-row independent, or masked. The MLX backend therefore omits it, which is a compute saving rather than a behavioural difference.

Checked against the reference rather than assumed. With a context of 100 rows (torch pads to 128), max |delta| on the real-row logits, which are on a 0.363 scale:

comparison	max abs diff
torch padding internally vs the caller pre-padding to 128	0.0
same, but the padded rows filled with garbage features at 50x scale	0.0
torch (padded) vs MLX (unpadded)	5.96e-07

Arbitrary values in the padded rows leave the real rows bit-identical, so the padding cannot carry information into the result. Two tests keep this honest rather than leaving it as an argument in a PR description: test_reference_padding_is_inert asserts the property on the reference itself (assert_array_equal, since a masked row may not perturb the result at all -- and it is not vacuous: swapping the sentinel labels for valid ones moves the logits by 2.8e-02), and test_unpadded_matches_torch_at_awkward_lengths checks the two backends agree at (100, 100), (100, 20), (5, 100), (1, 1) and (127, 129) train/test lengths.

Plus a 5-test opt-in release gate, real_checkpoint_parity_test.py, which runs the shipped estimators against PyTorch CPU in float32 — the reference both bfloat16 backends approximate — on the real v1.0.0 checkpoint. It covers the uncached forward, the cached prefill/decode path and the int8 KV cache. It's gated because it pulls the ~6 GB checkpoint and holds two model copies (~13 GB RSS):

TABFM_REAL_PARITY=1 python -m pytest tabfm/src/mlx/real_checkpoint_parity_test.py

The small-random-model tests pin the wiring but not the released weights — a transposed projection or a mis-set config key can pass there and still be wrong on the real checkpoint, which is what the gate is for.

Benchmark (Apple M3 Pro, same weights and data on both backends)

100 train / 20 test, interleaved A/B trials (running both at once would contend
for the same memory bandwidth), median of 3, seconds:

path	op	MLX	PyTorch/MPS	speedup
classification, uncached	predict	2.90	6.07	2.09x
regression, uncached	predict	2.91	6.10	2.09x
classification, cached	fit (prefill)	3.38	6.40	1.89x
classification, cached	predict (decode)	0.47	4.81	10.22x
regression, cached	fit (prefill)	3.39	5.94	1.75x
regression, cached	predict (decode)	0.47	4.66	9.90x

The gap is widest on the cached decode path, which is where the in-context
learning loop spends its time when predicting repeatedly against one fitted
context.

Agreement measured in the same runs: 100% class agreement with max
|delta proba| 1.4e-03 uncached and 4.3e-03 cached, and regression max |delta|
0.074 to 0.103 on a ~49.8 scale -- both backends running bfloat16.

On dtype

bfloat16 is the default and the model's design dtype. float32 costs ~1.2x
time and 8.47 GB peak for a max difference of 2.7e-07 from the PyTorch
float32 reference. float16 is rejected by the loader: the activations exceed
its 65504 range, the forward pass returns NaN and class agreement falls to
chance, and it is not faster than bfloat16 anyway (both 0.925s) because the
cost here is compute, not weight bandwidth.

On the norm kernel

RMSNorm uses MLX's fused mx.fast.rms_norm rather than a manual float32
chain. The 257 norm modules are ~23% of inference time by ablation, and the
fused kernel recovers essentially all of it. In float32 the two agree to
3e-07, so the parity gate above is unaffected. In bfloat16 the kernel rounds
x*rsqrt before applying the weight, but measured against the float32
reference the manual chain is 4.415e-03 out and the fused kernel 4.492e-03:
the bfloat16 path is already that far from float32, and predicted classes are
identical. mx.compile over the manual chain is the bit-identical
alternative at 1.12x if you would rather have that; the remaining levers were each measured and rejected: compiling the
MLP or the attention block (matmul-bound, no change), 8-bit weight
quantization (10% slower, since runtime scales linearly with ensemble size
and is therefore compute- not bandwidth-bound), mx.set_wired_limit (no
effect), and eliminating the attention transposes (5 ms of 2890 ms, since
MLX folds them into the consuming kernel).

The benchmark driver lives outside the repo; happy to add it under the MLX
directory if that is useful to you.

The loader's docstring notes that the default pretrained weights remain under the separate tabfm-non-commercial-v1.0 license.

[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.