Make fitted estimators picklable after predict (JAX backend)

[fus3r]

Summary

Saving a fitted estimator with stdlib pickle crashes on the JAX backend once predict has run:

AttributeError: Can't pickle local object 'TabFMClassifier._batch_forward.<locals>._predict_step_fn'

The first predict memoizes nnx.jit-compiled step functions on the estimator instance (_predict_step_compiled_with_cat / _predict_step_compiled_no_cat, see _batch_forward). Those closures cannot be pickled. This is the same save path that motivated #47 on the PyTorch side: AutoGluon / TabArena pickle the fitted estimator, and fit then predict then save is the normal order, so this breaks the JAX backend for that workflow.

The fix drops the memoized functions from __getstate__ on both estimators. They are pure caches, guarded by hasattr in _batch_forward, so they are rebuilt lazily on the next predict; a restored estimator produces identical predictions (covered by the new tests). dict(super().__getstate__()) is used because sklearn's BaseEstimator.__getstate__ can hand back the live __dict__ for estimators outside the sklearn namespace, and pickling should not evict the running object's cache as a side effect.

Pickling a fitted estimator before the first predict already worked; only the after-predict state was broken.

Repro on main

import pickle
import numpy as np
from flax import nnx
from tabfm.src.jax import model as tabfm_model
from tabfm.src.classifier_and_regressor import TabFMClassifier

model = tabfm_model.TabFM(
    loss="cross_entropy", max_classes=10, embed_dim=8, col_num_blocks=1,
    col_nhead=2, col_num_inds=8, row_num_blocks=1, row_nhead=2,
    row_num_cls=1, icl_num_blocks=1, icl_nhead=2, rngs=nnx.Rngs(0))
clf = TabFMClassifier(model=model, n_estimators=2)
X = np.random.rand(16, 3)
clf.fit(X, np.array([0, 1] * 8))
pickle.dumps(clf)          # works
clf.predict_proba(X[:4])
pickle.dumps(clf)          # AttributeError: Can't pickle local object ...

Testing

• 2 new tests (EstimatorPickleTest): fit, predict, pickle round trip, and the restored estimator's predictions match exactly (classifier and regressor).
• Full pytest tabfm/src/ passes (74 tests, both backends installed).