keep compressed_axes if it's consistent

pydata · Jul 2, 2023 · 683b2b8 · 683b2b8
1 parent 9a7e7d1
commit 683b2b8
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/sparse/_umath.py b/sparse/_umath.py
@@ -436,7 +436,8 @@ def __init__(self, func, *args, **kwargs):
             out_type = DOK
         elif all(isinstance(arg, GCXS) for arg in sparse_args):
             out_type = GCXS
-            out_kwargs["compressed_axes"] = sparse_args[0].compressed_axes
+            if len({arg.compressed_axes for arg in sparse_args}) == 1:
+                out_kwargs["compressed_axes"] = sparse_args[0].compressed_axes
         else:
             out_type = COO