fix(distributions): make distribution families picklable

deeptab/distributions/base.py

@@ -5,6 +5,26 @@
 import torch
 
 
+# Parameter transforms are module-level functions rather than lambdas so that
+# distribution families (and the estimators and Lightning modules that hold
+# them) can be pickled, e.g. for pickle.dumps() or the ddp_spawn strategy.
+def _identity(x: torch.Tensor) -> torch.Tensor:
+    return x
+
+
+def _square(x: torch.Tensor) -> torch.Tensor:
+    return x**2
+
+
+def _softmax_last_dim(x: torch.Tensor) -> torch.Tensor:
+    return torch.softmax(x, dim=-1)
+
+
+def _stable_log(x: torch.Tensor) -> torch.Tensor:
+    # Adding a small constant for numerical stability
+    return torch.log(x + 1e-6)
+
+
 class BaseDistribution(torch.nn.Module):
     """
     The base class for various statistical distributions, providing a common interface and utilities.
@@ -35,14 +55,13 @@ def __init__(self, name, param_names):
         # Predefined transformation functions accessible to all subclasses
         self.predefined_transforms: dict[str, Callable[[torch.Tensor], torch.Tensor]] = {
             "positive": torch.nn.functional.softplus,
-            "none": lambda x: x,
-            "square": lambda x: x**2,
+            "none": _identity,
+            "square": _square,
             "exp": torch.exp,
             "sqrt": torch.sqrt,
             "sigmoid": torch.sigmoid,
-            "probabilities": lambda x: torch.softmax(x, dim=-1),
-            # Adding a small constant for numerical stability
-            "log": lambda x: torch.log(x + 1e-6),
+            "probabilities": _softmax_last_dim,
+            "log": _stable_log,
         }
 
     @property
@@ -63,7 +82,7 @@ def get_transform(
             # Custom transformation function provided
             return transform_name
         # Default to 'none'
-        return self.predefined_transforms.get(transform_name, lambda x: x)
+        return self.predefined_transforms.get(transform_name, _identity)
 
     def compute_loss(self, predictions, y_true):
         """

tests/integration/test_save_load.py

@@ -224,6 +224,43 @@ def reject_task_serialization(self):
     np.testing.assert_array_equal(model.predict(features), expected)
 
 
+def _registered_families():
+    from deeptab.distributions import DISTRIBUTION_REGISTRY
+
+    return sorted(DISTRIBUTION_REGISTRY)
+
+
+@pytest.mark.parametrize("family", _registered_families())
+def test_every_distribution_family_survives_a_pickle_round_trip(family):
+    """Families hold their parameter transforms, which therefore must be picklable."""
+    from deeptab.distributions import get_distribution
+
+    distribution = get_distribution(family)
+    restored = pickle.loads(pickle.dumps(distribution))
+    assert type(restored) is type(distribution)
+    assert restored.param_names == distribution.param_names
+    if family == "normal":
+        raw = torch.randn(5, distribution.parameter_count)
+        torch.testing.assert_close(restored(raw), distribution(raw))
+
+
+@pytest.mark.parametrize("family", ["normal", "quantile", "categorical"])
+def test_fitted_lss_can_be_pickled(family, regression_data, tmp_path):
+    features, _, targets, _ = regression_data
+    if family == "categorical":
+        targets = (targets > np.median(targets)).astype(int)
+    model = MLPLSS(
+        model_config=MLPConfig(layer_sizes=[16]),
+        trainer_config=TrainerConfig(max_epochs=1, batch_size=64, checkpoint_path=str(tmp_path)),
+        random_state=42,
+    )
+    model.fit(features, targets, family=family, accelerator="cpu", logger=False, enable_progress_bar=False)
+    restored = pickle.loads(pickle.dumps(model))
+    assert restored.family_name == model.family_name
+    with pytest.raises(NotFittedError):
+        check_is_fitted(restored)
+
+
 def test_regressor_save_raises_when_unfitted():
     model = MLPRegressor()
     with pytest.raises(ValueError, match="fitted"):