ENH: Updates to numpy.array_api (#19937)

* Add __index__ to array_api and update __int__, __bool__, and __float__

The spec specifies that they should only work on arrays with corresponding
dtypes. __index__ is new in the spec since the initial PR, and works
identically to np.array.__index__.

* Add the to_device method to the array_api

This method is new since #18585. It does nothing in NumPy since NumPy does not
support non-CPU devices.

* Update transpose methods in the array_api

transpose() was renamed to matrix_transpose() and now operates on stacks of
matrices. A function to permute dimensions will be added once it is finalized
in the spec. The attribute mT was added and the T attribute was updated to
only operate on 2-dimensional arrays as per the spec.

* Restrict input dtypes in the array API statistical functions

* Add the dtype parameter to the array API sum() and prod()

* Add the function permute_dims() to the array_api namespace

permute_dims() is the replacement for transpose(), which was split into
permute_dims() and matrix_transpose().

* Add tril and triu to the array API namespace

* Fix the array_api Array.__repr__ to indent the array properly

* Make the Device type in the array_api just accept the string "cpu"

1 files changed, 33 insertions, 1 deletions

diff --git a/numpy/array_api/_array_object.py b/numpy/array_api/_array_object.py
index 2d746e78b..830319e8c 100644
--- a/numpy/array_api/_array_object.py
+++ b/numpy/array_api/_array_object.py
@@ -99,7 +99,10 @@ class Array:
         """
         Performs the operation __repr__.
         """
-        return f"Array({np.array2string(self._array, separator=', ')}, dtype={self.dtype.name})"
+        prefix = "Array("
+        suffix = f", dtype={self.dtype.name})"
+        mid = np.array2string(self._array, separator=', ', prefix=prefix, suffix=suffix)
+        return prefix + mid + suffix
 
     # These are various helper functions to make the array behavior match the
     # spec in places where it either deviates from or is more strict than
@@ -391,6 +394,8 @@ class Array:
         # Note: This is an error here.
         if self._array.ndim != 0:
             raise TypeError("bool is only allowed on arrays with 0 dimensions")
+        if self.dtype not in _boolean_dtypes:
+            raise ValueError("bool is only allowed on boolean arrays")
         res = self._array.__bool__()
         return res
 
@@ -429,6 +434,8 @@ class Array:
         # Note: This is an error here.
         if self._array.ndim != 0:
             raise TypeError("float is only allowed on arrays with 0 dimensions")
+        if self.dtype not in _floating_dtypes:
+            raise ValueError("float is only allowed on floating-point arrays")
         res = self._array.__float__()
         return res
 
@@ -488,9 +495,18 @@ class Array:
         # Note: This is an error here.
         if self._array.ndim != 0:
             raise TypeError("int is only allowed on arrays with 0 dimensions")
+        if self.dtype not in _integer_dtypes:
+            raise ValueError("int is only allowed on integer arrays")
         res = self._array.__int__()
         return res
 
+    def __index__(self: Array, /) -> int:
+        """
+        Performs the operation __index__.
+        """
+        res = self._array.__index__()
+        return res
+
     def __invert__(self: Array, /) -> Array:
         """
         Performs the operation __invert__.
@@ -979,6 +995,11 @@ class Array:
         res = self._array.__rxor__(other._array)
         return self.__class__._new(res)
 
+    def to_device(self: Array, device: Device, /) -> Array:
+        if device == 'cpu':
+            return self
+        raise ValueError(f"Unsupported device {device!r}")
+
     @property
     def dtype(self) -> Dtype:
         """
@@ -992,6 +1013,12 @@ class Array:
     def device(self) -> Device:
         return "cpu"
 
+    # Note: mT is new in array API spec (see matrix_transpose)
+    @property
+    def mT(self) -> Array:
+        from ._linear_algebra_functions import matrix_transpose
+        return matrix_transpose(self)
+
     @property
     def ndim(self) -> int:
         """
@@ -1026,4 +1053,9 @@ class Array:
 
         See its docstring for more information.
         """
+        # Note: T only works on 2-dimensional arrays. See the corresponding
+        # note in the specification:
+        # https://data-apis.org/array-api/latest/API_specification/array_object.html#t
+        if self.ndim != 2:
+            raise ValueError("x.T requires x to have 2 dimensions. Use x.mT to transpose stacks of matrices and permute_dims() to permute dimensions.")
         return self._array.T