ENH: Add the linalg extension to the array_api submodule (numpy#19980)

Author: asmeurer

numpy/array_api/__init__.py

@@ -109,9 +109,6 @@
 - The spec is still in an RFC phase and may still have minor updates, which
   will need to be reflected here.
 
-- The linear algebra extension in the spec will be added in a future pull
-  request.
-
 - Complex number support in array API spec is planned but not yet finalized,
   as are the fft extension and certain linear algebra functions such as eig
   that require complex dtypes.
@@ -334,12 +331,13 @@
     "trunc",
 ]
 
-# einsum is not yet implemented in the array API spec.
+# linalg is an extension in the array API spec, which is a sub-namespace. Only
+# a subset of functions in it are imported into the top-level namespace.
+from . import linalg
 
-# from ._linear_algebra_functions import einsum
-# __all__ += ['einsum']
+__all__ += ["linalg"]
 
-from ._linear_algebra_functions import matmul, tensordot, matrix_transpose, vecdot
+from .linalg import matmul, tensordot, matrix_transpose, vecdot
 
 __all__ += ["matmul", "tensordot", "matrix_transpose", "vecdot"]
 

numpy/array_api/_array_object.py

@@ -1030,7 +1030,7 @@ def device(self) -> Device:
     # Note: mT is new in array API spec (see matrix_transpose)
     @property
     def mT(self) -> Array:
-        from ._linear_algebra_functions import matrix_transpose
+        from .linalg import matrix_transpose
         return matrix_transpose(self)
 
     @property

numpy/array_api/_linear_algebra_functions.py

@@ -93,11 +93,12 @@ def sum(
 ) -> Array:
     if x.dtype not in _numeric_dtypes:
         raise TypeError("Only numeric dtypes are allowed in sum")
-    # Note: sum() and prod() always upcast float32 to float64 for dtype=None
-    # We need to do so here before summing to avoid overflow
+    # Note: sum() and prod() always upcast integers to (u)int64 and float32 to
+    # float64 for dtype=None. `np.sum` does that too for integers, but not for
+    # float32, so we need to special-case it here
     if dtype is None and x.dtype == float32:
-        x = asarray(x, dtype=float64)
-    return Array._new(np.sum(x._array, axis=axis, keepdims=keepdims))
+        dtype = float64
+    return Array._new(np.sum(x._array, axis=axis, dtype=dtype, keepdims=keepdims))
 
 
 def var(