Backport some fixes.

pydata · May 21, 2024 · 32d9f80 · 32d9f80
1 parent 34b7185
commit 32d9f80
Show file tree

Hide file tree

Showing 5 changed files with 82 additions and 82 deletions.
diff --git a/sparse/__init__.py b/sparse/__init__.py
@@ -23,7 +23,6 @@
     greater,
     greater_equal,
     iinfo,
-    imag,
     inf,
     int8,
     int16,
@@ -47,7 +46,6 @@
     not_equal,
     pi,
     positive,
-    real,
     remainder,
     sign,
     sin,
@@ -160,6 +158,9 @@
     "GCXS",
     "SparseArray",
     "abs",
+    "acos",
+    "acosh",
+    "add",
     "all",
     "any",
     "argmax",
@@ -168,22 +169,51 @@
     "asCOO",
     "as_coo",
     "asarray",
+    "asin",
+    "asinh",
     "asnumpy",
     "astype",
+    "atan",
+    "atan2",
+    "atanh",
+    "bitwise_and",
+    "bitwise_invert",
+    "bitwise_left_shift",
+    "bitwise_not",
+    "bitwise_or",
+    "bitwise_right_shift",
+    "bitwise_xor",
+    "bool",
     "broadcast_arrays",
     "broadcast_to",
+    "can_cast",
+    "ceil",
     "clip",
+    "complex128",
+    "complex64",
     "concat",
     "concatenate",
+    "cos",
+    "cosh",
     "diagonal",
     "diagonalize",
+    "divide",
     "dot",
+    "e",
     "einsum",
     "elemwise",
     "empty",
     "empty_like",
     "equal",
+    "exp",
+    "expm1",
     "eye",
+    "finfo",
+    "float16",
+    "float32",
+    "float64",
+    "floor",
+    "floor_divide",
     "full",
     "full_like",
     "greater",
@@ -201,25 +231,43 @@
     "isneginf",
     "isposinf",
     "kron",
+    "less",
+    "less_equal",
     "load_npz",
+    "log",
+    "log10",
+    "log1p",
+    "log2",
+    "logaddexp",
+    "logical_and",
+    "logical_not",
+    "logical_or",
+    "logical_xor",
     "matmul",
     "max",
     "mean",
     "min",
     "moveaxis",
-    "moveaxis",
+    "multiply",
+    "nan",
     "nanmax",
     "nanmean",
     "nanmin",
     "nanprod",
     "nanreduce",
     "nansum",
+    "negative",
+    "newaxis",
     "nonzero",
+    "not_equal",
     "ones",
     "ones_like",
     "outer",
     "pad",
     "permute_dims",
+    "pi",
+    "positive",
+    "pow",
     "prod",
     "random",
     "real",
@@ -229,84 +277,31 @@
     "roll",
     "round",
     "save_npz",
-    "squeeze",
-    "stack",
-    "std",
-    "sum",
-    "tensordot",
-    "tril",
-    "triu",
-    "unique_counts",
-    "unique_values",
-    "where",
-    "zeros",
-    "zeros_like",
-    "add",
-    "bitwise_and",
-    "bitwise_not",
-    "bitwise_or",
-    "bitwise_xor",
-    "can_cast",
-    "ceil",
-    "complex64",
-    "complex128",
-    "cos",
-    "cosh",
-    "divide",
-    "e",
-    "exp",
-    "expm1",
-    "finfo",
-    "float16",
-    "float32",
-    "float64",
-    "floor",
-    "floor_divide",
-    "less",
-    "less_equal",
-    "log",
-    "log1p",
-    "log2",
-    "log10",
-    "logaddexp",
-    "logical_and",
-    "logical_not",
-    "logical_or",
-    "logical_xor",
-    "multiply",
-    "nan",
-    "negative",
-    "newaxis",
-    "not_equal",
-    "pi",
-    "positive",
-    "real",
-    "remainder",
     "sign",
     "sin",
     "sinh",
     "sqrt",
     "square",
+    "squeeze",
+    "stack",
+    "std",
     "subtract",
+    "sum",
     "tan",
     "tanh",
+    "tensordot",
+    "tril",
+    "triu",
     "trunc",
-    "uint8",
     "uint16",
     "uint32",
     "uint64",
-    "acos",
-    "acosh",
-    "asin",
-    "asinh",
-    "atan",
-    "atan2",
-    "atanh",
-    "bool",
-    "bitwise_invert",
-    "bitwise_left_shift",
-    "pow",
-    "bitwise_right_shift",
+    "uint8",
+    "unique_counts",
+    "unique_values",
+    "where",
+    "zeros",
+    "zeros_like",
 ]
 
 __array_api_version__ = "2022.12"
diff --git a/sparse/_common.py b/sparse/_common.py
@@ -2188,6 +2188,16 @@ def isnan(x, /):
     return x.isnan()
 
 
+@_support_numpy
+def imag(x, /):
+    return x.imag
+
+
+@_support_numpy
+def real(x, /):
+    return x.real
+
+
 def isfinite(x, /):
     return ~isinf(x)
 

diff --git a/sparse/_compressed/compressed.py b/sparse/_compressed/compressed.py
@@ -319,10 +319,10 @@ def T(self):
         return self.transpose()
 
     def __str__(self):
-        summary = "<GCXS: shape={}, dtype={}, nnz={}, fill_value={}, compressed_axes={}>".format(
-            self.shape, self.dtype, self.nnz, self.fill_value, self.compressed_axes
+        return self._str_impl(
+            f"<GCXS: shape={self.shape}, dtype={self.dtype}, nnz={self.nnz}, "
+            f"fill_value={self.fill_value}, compressed_axes={self.compressed_axes}>"
         )
-        return self._str_impl(summary)
 
     __repr__ = __str__
 
@@ -832,14 +832,10 @@ def __init__(self, arg, shape=None, compressed_axes=None, prune=False, fill_valu
         )
 
     def __str__(self):
-        summary = "<{}: shape={}, dtype={}, nnz={}, fill_value={}>".format(
-            type(self).__name__,
-            self.shape,
-            self.dtype,
-            self.nnz,
-            self.fill_value,
+        return self._str_impl(
+            f"<{type(self).__name__}: shape={self.shape}, dtype={self.dtype}, "
+            f"nnz={self.nnz}, fill_value={self.fill_value}>"
         )
-        return self._str_impl(summary)
 
     __repr__ = __str__
 

diff --git a/sparse/_coo/common.py b/sparse/_coo/common.py
@@ -229,11 +229,9 @@ def stack(arrays, axis=0):
     shape.insert(axis, len(arrays))
 
     nnz = 0
-    dim = 0
     new = np.empty(shape=(coords.shape[1],), dtype=np.intp)
-    for x in arrays:
+    for dim, x in enumerate(arrays):
         new[nnz : x.nnz + nnz] = dim
-        dim += 1
         nnz += x.nnz
 
     coords = [coords[i] for i in range(coords.shape[0])]

diff --git a/sparse/_coo/numba_extension.py b/sparse/_coo/numba_extension.py
@@ -3,6 +3,7 @@
 
 For now, this just supports attribute access
 """
+
 import contextlib
 
 import numba