Using np.asarray instead

simonsobs · Feb 16, 2024 · 3c063d7 · 3c063d7
1 parent 603029d
commit 3c063d7
Showing 1 changed file with 4 additions and 8 deletions.
diff --git a/pixell/utils.py b/pixell/utils.py
@@ -1236,16 +1236,12 @@ def pole_wrap(pos):
 def allreduce(a, comm, op=None):
     """Convenience wrapper for Allreduce that returns the result
     rather than needing an output argument."""
-    data_to_reduce = a.copy()
-    if isinstance(a, np.generic):
-        data_to_reduce = np.array([data_to_reduce])
-    res = np.zeros_like(data_to_reduce)
+    a = np.asarray(a)
+    res = np.zeros_like(a)
     if op is None:
-        comm.Allreduce(data_to_reduce, res)
+        comm.Allreduce(a, res)
     else:
-        comm.Allreduce(data_to_reduce, res, op)
-    if isinstance(a, np.generic):
-        res = res[0]
+        comm.Allreduce(a, res, op)
     return res
 
 def reduce(a, comm, root=0, op=None):