keep compressed_axes if it's consistent

pydata · Jan 5, 2024 · 0d04116 · 0d04116
1 parent 1311e4c
commit 0d04116
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/sparse/_umath.py b/sparse/_umath.py
@@ -422,7 +422,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