Ev: Add map-reduce style interface.

src/gluonts/ev/aggregations.py

@@ -11,11 +11,16 @@
 # express or implied. See the License for the specific language governing
 # permissions and limitations under the License.
 
+from __future__ import annotations
+
+
 from dataclasses import dataclass
 from typing import List, Optional, Union
 
 import numpy as np
 
+from gluonts import maybe
+
 
 @dataclass
 class Aggregation:
@@ -31,6 +36,9 @@ def step(self, values: np.ndarray) -> None:
     def get(self) -> np.ndarray:
         raise NotImplementedError
 
+    def reduce(self, other) -> Aggregation:
+        raise NotImplementedError
+
 
 @dataclass
 class Sum(Aggregation):
@@ -65,6 +73,15 @@ def get(self) -> np.ndarray:
 
         return np.ma.concatenate(self.partial_result)
 
+    def reduce(self, other) -> Sum:
+        if self.axis is None or 0 in self.axis:
+            return Sum(self.axis, self.partial_result + other.partial_result)
+
+        return Sum(
+            self.axis,
+            np.concatenate([self.partial_result, other.partial_result]),
+        )
+
 
 @dataclass
 class Mean(Aggregation):
@@ -108,3 +125,21 @@ def get(self) -> np.ndarray:
             return self.partial_result / self.n
 
         return np.ma.concatenate(self.partial_result)
+
+    def reduce(self, other) -> Sum:
+        if self.axis is None or 0 in self.axis:
+            return Mean(
+                self.axis,
+                self.partial_result + other.partial_result,
+                self.n + other.n,
+            )
+
+        return Mean(
+            self.axis,
+            np.concatenate(
+                [
+                    maybe.unwrap_or(self.partial_result, []),
+                    maybe.unwrap_or(other.partial_result, []),
+                ]
+            ),
+        )

src/gluonts/ev/metrics.py

@@ -28,6 +28,8 @@
 
 import numpy as np
 
+from gluonts.itertools import join_items
+
 from .aggregations import Aggregation, Mean, Sum
 from .stats import (
     error,
@@ -89,6 +91,9 @@ def update_all(self, stream: Iterator[Mapping[str, np.ndarray]]) -> Self:
     def get(self) -> np.ndarray:
         raise NotImplementedError
 
+    def reduce(self, other):
+        raise NotImplementedError
+
 
 @dataclass
 class DirectMetric(Metric):
@@ -105,6 +110,11 @@ def update(self, data: Mapping[str, np.ndarray]) -> Self:
     def get(self) -> np.ndarray:
         return self.aggregate.get()
 
+    def reduce(self, other):
+        return DirectMetric(
+            self.name, self.stat, self.aggregate.reduce(other.aggregate)
+        )
+
 
 @dataclass
 class DerivedMetric(Metric):
@@ -131,6 +141,20 @@ def get(self) -> np.ndarray:
             }
         )
 
+    def reduce(self, other):
+        metrics = {
+            key: left.reduce(right)
+            for key, left, right in join_items(
+                self.metrics, other.metrics, "strict"
+            )
+        }
+
+        return DerivedMetric(
+            self.name,
+            metrics,
+            self.post_process,
+        )
+
 
 @runtime_checkable
 class MetricDefinition(Protocol):

src/gluonts/torch/__init__.py

@@ -24,6 +24,7 @@
     "LagTSTEstimator",
 ]
 import torch
+<<<<<<< Updated upstream
 from .model.estimator import PyTorchLightningEstimator
 from .model.predictor import PyTorchPredictor
 from .model.deep_npts import DeepNPTSEstimator
@@ -34,6 +35,24 @@
 from .model.d_linear import DLinearEstimator
 from .model.patch_tst import PatchTSTEstimator
 from .model.lag_tst import LagTSTEstimator
+=======
+
+from gluonts.meta.export import re_export
+
+__all__ = re_export(
+    __name__,
+    "model.estimator:PyTorchLightningEstimator",
+    "model.predictor:PyTorchPredictor",
+    "model.deep_npts:DeepNPTSEstimator",
+    "model.deepar:DeepAREstimator",
+    "model.simple_feedforward:SimpleFeedForwardEstimator",
+    "model.tft:TemporalFusionTransformerEstimator",
+    "model.wavenet:WaveNetEstimator",
+    "model.d_linear:DLinearEstimator",
+    "model.patch_tst:PatchTSTEstimator",
+    "model.lag_tst:LagTSTEstimator",
+)
+>>>>>>> Stashed changes
 
 
 from . import prelude as _  # noqa