Fix metric per rank computation

philipphager · philipphager · commit eede2925a3b8 · 2025-10-31T15:20:18.000-04:00
diff --git a/clax/metrics.py b/clax/metrics.py
@@ -46,7 +46,7 @@ def compute(self, prefix: str = "") -> Dict[str, Any]:
 
     def compute_per_rank(self, prefix: str = "") -> Dict[str, Any]:
         return {
-            f"{prefix}{name}": getattr(self, name).compute_per_rank()
+            f"{prefix}{name}": metric.compute_per_rank()
             for name, metric in self.metrics.items()
             if isinstance(metric, RankBasedAverage)
         }

Original file line number	Diff line number	Diff line change
`@@ -46,7 +46,7 @@ def compute(self, prefix: str = "") -> Dict[str, Any]:`
`46`	`46`
`47`	`47`	`def compute_per_rank(self, prefix: str = "") -> Dict[str, Any]:`
`48`	`48`	`return {`
`49`		`- f"{prefix}{name}": getattr(self, name).compute_per_rank()`
	`49`	`+ f"{prefix}{name}": metric.compute_per_rank()`
`50`	`50`	`for name, metric in self.metrics.items()`
`51`	`51`	`if isinstance(metric, RankBasedAverage)`
`52`	`52`	`}`