Extend save/load serialization to non-meta learners (causalml.inference.*)

[jeongyoonlee]

Summary

#929 (closes #892) added save()/load() + a generic load_learner() to the meta-learners via a SerializableLearner mixin on causalml.inference.meta.base.BaseLearner, with joblib persistence and version/class metadata guardrails. This issue tracks extending that capability to the rest of causalml.inference.* (uplift trees, causal trees/forest, IV, and the neural-net learners), and records the design decisions that extension requires.

Two generalizations the mixin needs first

The mixin is generic in its useful parts (joblib+metadata wrapper, version/class checks, load_learner), but it's coupled to meta-learners in two places:

1. Fitted-check is meta-specific. save() guards with hasattr(self, "t_groups") (serialization.py:99). t_groups is a meta/DRIV concept — uplift trees signal fitted-ness with fitted_uplift_tree / uplift_forest, causal trees/forest with sklearn's tree_ / estimators_. Replace with an overridable _is_fitted() hook (default to sklearn.utils.validation.check_is_fitted(self) for the sklearn-based estimators; override for uplift trees; keep t_groups for meta/DRIV).
2. Location. The mixin lives in inference/meta/serialization.py and is mixed into the meta base. To share it, lift it to a neutral module (e.g. causalml/inference/serialization.py) and mix it into each family's base class.

Per-family feasibility (the real gate is picklability)

Group A — joblib round-trips, low effort:

• CausalTreeRegressor, CausalRandomForestRegressor — sklearn estimators, joblib works natively.
• UpliftTreeClassifier, UpliftRandomForestClassifier — plain Python; pickle-safe on Cython 3.x (verified during the UpliftRandomForestClassifier Model object cannot be pickled when saving as joblib #588 investigation).
• IVRegressor, BaseDRIVLearner/BaseDRIVRegressor — plain Python; DRIV already sets self.t_groups (iv/drivlearner.py:116), so it works with the mixin nearly unchanged.

Note: the sklearn-based and uplift learners are already joblib-dumpable today, so for them the mixin mainly adds the version/metadata guardrails and a consistent API — not new capability.

Group B — neural nets, need framework-native serialization (do NOT use joblib):

• DragonNet (TF and JAX) wraps a Keras/flax model — a joblib dump is fragile/broken; it already exposes native save/load (Keras H5 / flax).
• CEVAE wraps torch modules — needs state_dict via torch.save.
• Keep these on their native serialization; if unified under the shared API, the method should delegate to the framework-native path, not joblib.

Naming / collision decision

save()/load() are already defined on the NN learners with a different contract than the mixin's:

• tf/dragonnet.py:295,304 — save(self, h5_filepath), load(self, h5_filepath, ratio=1.0, dragonnet_loss=...) (instance method, mutates self).
• jax/dragonnet.py:484,507 — save(self, path), load(self, path, input_dim=None).

The mixin's load() is a classmethod returning a new object. Standardizing the cross-family verb on save()/load() would collide with these (and load_model() is a poor choice — "model" already denotes the inner base estimator, and it clashes with Keras' own load_model).

Decision: standardize the shared cross-family verb on save_learner() / load_learner() — it matches causalml's own noun ("learner"), reuses the load_learner() name #929 already ships, and doesn't clobber the NN learners' native save/load. Meta-learners keep their current save()/load()/load_learner() surface (no collision there); no change needed to #929.

Suggested increments

• PR 1 (Group A): lift + generalize the mixin (overridable _is_fitted(), neutral module) and apply to causal tree, causal forest, uplift tree, uplift forest, IV/DRIV. Round-trip tests per family.
• PR 2 (Group B): unify the API for DragonNet/CEVAE by delegating to Keras/flax/torch native serialization; guard tests with importorskip.

Refs: #892, #929.