Match all-reduce op to query_method (AVG for mean, SUM for sum)

luciaquirke · claude · luciaquirke · commit 77497a670b62 · 2026-03-07T11:26:40.000+11:00
Co-Authored-By: Claude Opus 4.6 &lt;noreply@anthropic.com&gt;
diff --git a/bergson/double_backward.py b/bergson/double_backward.py
@@ -203,8 +203,9 @@ def schedule(step: Numeric) -> Numeric:
     )
 
     if world_size > 1:
+        reduce_op = dist.ReduceOp.AVG if run_cfg.query_method == "mean" else dist.ReduceOp.SUM
         for v in query_grads.values():
-            dist.all_reduce(v, op=dist.ReduceOp.AVG)
+            dist.all_reduce(v, op=reduce_op)
 
     stream.requires_grad = True
     opt_grads = [

Original file line number	Diff line number	Diff line change
`@@ -203,8 +203,9 @@ def schedule(step: Numeric) -> Numeric:`
`203`	`203`	`)`
`204`	`204`
`205`	`205`	`if world_size > 1:`
	`206`	`+ reduce_op = dist.ReduceOp.AVG if run_cfg.query_method == "mean" else dist.ReduceOp.SUM`
`206`	`207`	`for v in query_grads.values():`
`207`		`- dist.all_reduce(v, op=dist.ReduceOp.AVG)`
	`208`	`+ dist.all_reduce(v, op=reduce_op)`
`208`	`209`
`209`	`210`	`stream.requires_grad = True`
`210`	`211`	`opt_grads = [`