diff --git a/deeptab/models/lss_base.py b/deeptab/models/lss_base.py index 1b1c7365..7ec44b6c 100644 --- a/deeptab/models/lss_base.py +++ b/deeptab/models/lss_base.py @@ -28,6 +28,29 @@ from deeptab.training import TaskModel +def _as_numpy(value): + """Return *value* as a NumPy array; tensors are moved to the CPU first.""" + return value.detach().cpu().numpy() if isinstance(value, torch.Tensor) else np.asarray(value) + + +def _same_family_options(options: dict, stored: dict) -> bool: + """Compare family options, treating arrays and tensors (e.g. quantile levels) by value. + + A plain ``options != stored`` raises for array values, because NumPy and + torch cannot turn an element-wise comparison into a single truth value. + """ + if options.keys() != stored.keys(): + return False + for key, value in options.items(): + other = stored[key] + if isinstance(value, np.ndarray | torch.Tensor) or isinstance(other, np.ndarray | torch.Tensor): + if not np.array_equal(_as_numpy(value), _as_numpy(other)): + return False + elif value != other: + return False + return True + + class SklearnBaseLSS(SklearnBase): """Distributional regression base class (LSS variant of SklearnBase). @@ -59,7 +82,9 @@ def _configure_distribution(self, y, family=None, distributional_kwargs=None, *, raise ValueError("Dirichlet num_classes must match the number of target components") options["num_classes"] = count if not allow_change: - if family != getattr(self, "family_name", None) or options != getattr(self, "distributional_kwargs_", {}): + if family != getattr(self, "family_name", None) or not _same_family_options( + options, getattr(self, "distributional_kwargs_", {}) + ): raise ValueError("Changing the distribution family or options requires rebuild=True") return self.family = get_distribution(family, **options) diff --git a/tests/integration/test_lss_base.py b/tests/integration/test_lss_base.py index d9a45eb1..138a6297 100644 --- a/tests/integration/test_lss_base.py +++ b/tests/integration/test_lss_base.py @@ -532,6 +532,33 @@ def test_reused_build_rejects_changed_options_without_mutating_family(): assert model.family_name == "normal" +@pytest.mark.parametrize( + "quantiles", [np.array([0.1, 0.5, 0.9]), torch.tensor([0.1, 0.5, 0.9], dtype=torch.float64)], ids=["numpy", "torch"] +) +def test_reused_build_accepts_array_options(quantiles): + model = _small_lss() + model.build_model(_X[:24], _Y[:24], family="quantile", distributional_kwargs={"quantiles": quantiles}) + model.fit(_X[:24], _Y[:24], family="quantile", rebuild=False) + model.fit(_X[:24], _Y[:24], family="quantile", distributional_kwargs={"quantiles": [0.1, 0.5, 0.9]}, rebuild=False) + assert model.is_fitted_ + assert model.predict(_X[:4]).shape == (4, 3) + + +def test_reused_build_rejects_different_numpy_array_options(): + model = _small_lss() + model.build_model( + _X[:24], _Y[:24], family="quantile", distributional_kwargs={"quantiles": np.array([0.1, 0.5, 0.9])} + ) + with pytest.raises(ValueError, match="requires rebuild=True"): + model.fit( + _X[:24], + _Y[:24], + family="quantile", + distributional_kwargs={"quantiles": np.array([0.2, 0.5, 0.8])}, + rebuild=False, + ) + + def test_legacy_lss_bundle_loads_without_distributional_kwargs(tmp_path, fitted_mlplss): path = str(tmp_path / "legacy.deeptab") fitted_mlplss.save(path)