MLflow for Darts implementation

.gitignore

@@ -25,3 +25,6 @@ docs_env
 .venv
 .env
 uv.lock
+repl/
+mlruns/*
+examples/mlruns/*

CHANGELOG.md

@@ -11,6 +11,12 @@ but cannot always guarantee backwards compatibility. Changes that may **break co
 
 **Improved**
 
+- 🚀🚀 Added a custom MLflow model flavor for Darts, available under `darts.utils.mlflow`. It provides an MLflow integration for any Darts `ForecastingModel` (statistical, scikit-learn-like, and PyTorch-based). [#3022](https://github.com/unit8co/darts/pull/3022) by [Jakub Chłapek](https://github.com/jakubchlapek), [Zhihao Dai](https://github.com/daidahao) and [Michel Zeller](https://github.com/mizeller).
+  - Added `save_model()`, `load_model()`, and `log_model()` to persist and reload Darts models as MLflow models, including model and covariate metadata.
+  - Added `autolog()` to automatically log model creation parameters, covariate usage, metrics, and the trained model artifact on every `fit()`, `backtest()` and `historical_forecasts()` call within an active MLflow run.
+  - Metric functions from `darts.metrics` are automatically logged when called inside an active run, with keys reflecting the metric output shape (per-component, per-quantile/interval, per-label, and per-timestep results charted across MLflow steps). Multi-series results log the mean over series and write the full per-series breakdown to a CSV artifact.
+  - Added a new notebook for [MLflow quickstart](https://unit8co.github.io/darts/examples/29-MLflow-quickstart.html) with detailed usage examples.
+  - Note: model serving and deployment (MLflow's `pyfunc` flavor, model signatures, and input examples) are not yet supported.
 - 🔴 Improved `TransformerModel` with proper encoder-decoder transformer architecture using teacher forcing during training and autoregressive inference, aligning the implementation with [Vaswani et al., 2017](https://arxiv.org/abs/1706.03762). Previously trained checkpoints are incompatible and must be retrained. [#1915](https://github.com/unit8co/darts/pull/1915) by [Jan Fidor](https://github.com/JanFidor) and [Dennis Bader](https://github.com/dennisbader).
 - Added a `CITATION.cff` file with the recommended citation metadata for Darts. [#3147](https://github.com/unit8co/darts/pull/3147) by [Zhihao Dai](https://github.com/daidahao).
 
@@ -20,6 +26,8 @@ but cannot always guarantee backwards compatibility. Changes that may **break co
 
 **Dependencies**
 
+- Added an optional `mlflow` extra (`mlflow>=3.0`) enabling the MLflow integration. [#3022](https://github.com/unit8co/darts/pull/3022) by [Jakub Chłapek](https://github.com/jakubchlapek), [Zhihao Dai](https://github.com/daidahao) and [Michel Zeller](https://github.com/mizeller).
+
 ### For developers of the library:
 
 - Simplified internal error handling: removed `raise_if` and `raise_if_not` in favor of `raise_log`, and made `raise_log` automatically resolve the caller's logger (no longer requires passing `logger` explicitly). [#3126](https://github.com/unit8co/darts/pull/3126) by [Dennis Bader](https://github.com/dennisbader).

darts/models/forecasting/pl_forecasting_module.py

@@ -450,8 +450,11 @@ def _update_metrics(self, output, target, metrics):
             pred = output.squeeze(dim=-1)
 
         # torch metrics require 2D targets of shape (batch size * ocl, num targets)
-        target = target.reshape(-1, self.n_targets)
-        pred = pred.reshape(-1, self.n_targets)
+        # contiguous() is needed because model outputs can be non-contiguous views
+        # (e.g. NBEATS slices the last dimension), and some torchmetrics implementations
+        # call .view() internally which requires a contiguous tensor.
+        target = target.reshape(-1, self.n_targets).contiguous()
+        pred = pred.reshape(-1, self.n_targets).contiguous()
 
         metrics.update(pred, target)
 

darts/tests/conftest.py

@@ -32,6 +32,7 @@ def _package_available(*names: str) -> bool:
 PLOTLY_AVAILABLE = _package_available("plotly")
 IPYTHON_AVAILABLE = _package_available("IPython")
 TIREX_AVAILABLE = _package_available("tirex")
+MLFLOW_AVAILABLE = _package_available("mlflow")
 
 tfm_kwargs: dict[str, Any] = {
     "pl_trainer_kwargs": {