Make the PyTorch model picklable (module-level gelu activation)

[weihaokong]

Problem

Pickling the PyTorch TabFM model fails:

AttributeError: Can't pickle local object 'get_activation.<locals>.<lambda>'

get_activation("gelu") returns a lambda, which is stored as self.act on the gelu MLP heads. The transformer FFNs use swiglu (F.silu, picklable), but every model carries at least one gelu MLP — the y-encoder and decoder (the regression path also adds a y-embedder) — so the model as a whole is unpicklable.

This matters because AutoGluon / TabArena save the fitted estimator (which holds the model) with stdlib pickle. TabArena currently works around it with a custom __getstate__/__setstate__ that drops and reloads the network (autogluon/tabarena#434); this fix lets that workaround be removed.

Repro on main:

import pickle
from tabfm.src.pytorch import model
m = 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)
pickle.dumps(m)  # AttributeError

Fix

Promote the gelu activation to a module-level _gelu_tanh function. Module-level functions pickle by reference; the computation is unchanged (F.gelu(x, approximate="tanh")), so weights and outputs are identical — this is purely a serialization fix.

Testing

Adds PyTorchModelPickleTest (torch-only): classifier and regressor model pickle round-trips, plus a forward-output equivalence check after unpickling. The regressor case exercises all three gelu MLP heads.

pytest tabfm/src/classifier_and_regressor_pytorch_test.py passes (5 tests).