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):