1 files changed, 129 insertions, 29 deletions
diff --git a/numpy/array_api/_creation_functions.py b/numpy/array_api/_creation_functions.py
index acf78056a..e9c01e7e6 100644
--- a/numpy/array_api/_creation_functions.py
+++ b/numpy/array_api/_creation_functions.py
@@ -2,14 +2,22 @@ from __future__ import annotations
 
 
 from typing import TYPE_CHECKING, List, Optional, Tuple, Union
+
 if TYPE_CHECKING:
-    from ._typing import (Array, Device, Dtype, NestedSequence,
-                          SupportsDLPack, SupportsBufferProtocol)
+    from ._typing import (
+        Array,
+        Device,
+        Dtype,
+        NestedSequence,
+        SupportsDLPack,
+        SupportsBufferProtocol,
+    )
     from collections.abc import Sequence
 from ._dtypes import _all_dtypes
 
 import numpy as np
 
+
 def _check_valid_dtype(dtype):
     # Note: Only spelling dtypes as the dtype objects is supported.
 
@@ -20,7 +28,23 @@ def _check_valid_dtype(dtype):
             return
     raise ValueError("dtype must be one of the supported dtypes")
 
-def asarray(obj: Union[Array, bool, int, float, NestedSequence[bool|int|float], SupportsDLPack, SupportsBufferProtocol], /, *, dtype: Optional[Dtype] = None, device: Optional[Device] = None, copy: Optional[bool] = None) -> Array:
+
+def asarray(
+    obj: Union[
+        Array,
+        bool,
+        int,
+        float,
+        NestedSequence[bool | int | float],
+        SupportsDLPack,
+        SupportsBufferProtocol,
+    ],
+    /,
+    *,
+    dtype: Optional[Dtype] = None,
+    device: Optional[Device] = None,
+    copy: Optional[bool] = None,
+) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.asarray <numpy.asarray>`.
 
@@ -31,7 +55,7 @@ def asarray(obj: Union[Array, bool, int, float, NestedSequence[bool|int|float],
     from ._array_object import Array
 
     _check_valid_dtype(dtype)
-    if device not in ['cpu', None]:
+    if device not in ["cpu", None]:
         raise ValueError(f"Unsupported device {device!r}")
     if copy is False:
         # Note: copy=False is not yet implemented in np.asarray
@@ -40,14 +64,23 @@ def asarray(obj: Union[Array, bool, int, float, NestedSequence[bool|int|float],
         if copy is True:
             return Array._new(np.array(obj._array, copy=True, dtype=dtype))
         return obj
-    if dtype is None and isinstance(obj, int) and (obj > 2**64 or obj < -2**63):
+    if dtype is None and isinstance(obj, int) and (obj > 2 ** 64 or obj < -(2 ** 63)):
         # Give a better error message in this case. NumPy would convert this
         # to an object array. TODO: This won't handle large integers in lists.
         raise OverflowError("Integer out of bounds for array dtypes")
     res = np.asarray(obj, dtype=dtype)
     return Array._new(res)
 
-def arange(start: Union[int, float], /, stop: Optional[Union[int, float]] = None, step: Union[int, float] = 1, *, dtype: Optional[Dtype] = None, device: Optional[Device] = None) -> Array:
+
+def arange(
+    start: Union[int, float],
+    /,
+    stop: Optional[Union[int, float]] = None,
+    step: Union[int, float] = 1,
+    *,
+    dtype: Optional[Dtype] = None,
+    device: Optional[Device] = None,
+) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.arange <numpy.arange>`.
 
@@ -56,11 +89,17 @@ def arange(start: Union[int, float], /, stop: Optional[Union[int, float]] = None
     from ._array_object import Array
 
     _check_valid_dtype(dtype)
-    if device not in ['cpu', None]:
+    if device not in ["cpu", None]:
         raise ValueError(f"Unsupported device {device!r}")
     return Array._new(np.arange(start, stop=stop, step=step, dtype=dtype))
 
-def empty(shape: Union[int, Tuple[int, ...]], *, dtype: Optional[Dtype] = None, device: Optional[Device] = None) -> Array:
+
+def empty(
+    shape: Union[int, Tuple[int, ...]],
+    *,
+    dtype: Optional[Dtype] = None,
+    device: Optional[Device] = None,
+) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.empty <numpy.empty>`.
 
@@ -69,11 +108,14 @@ def empty(shape: Union[int, Tuple[int, ...]], *, dtype: Optional[Dtype] = None,
     from ._array_object import Array
 
     _check_valid_dtype(dtype)
-    if device not in ['cpu', None]:
+    if device not in ["cpu", None]:
         raise ValueError(f"Unsupported device {device!r}")
     return Array._new(np.empty(shape, dtype=dtype))
 
-def empty_like(x: Array, /, *, dtype: Optional[Dtype] = None, device: Optional[Device] = None) -> Array:
+
+def empty_like(
+    x: Array, /, *, dtype: Optional[Dtype] = None, device: Optional[Device] = None
+) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.empty_like <numpy.empty_like>`.
 
@@ -82,11 +124,20 @@ def empty_like(x: Array, /, *, dtype: Optional[Dtype] = None, device: Optional[D
     from ._array_object import Array
 
     _check_valid_dtype(dtype)
-    if device not in ['cpu', None]:
+    if device not in ["cpu", None]:
         raise ValueError(f"Unsupported device {device!r}")
     return Array._new(np.empty_like(x._array, dtype=dtype))
 
-def eye(n_rows: int, n_cols: Optional[int] = None, /, *, k: Optional[int] = 0, dtype: Optional[Dtype] = None, device: Optional[Device] = None) -> Array:
+
+def eye(
+    n_rows: int,
+    n_cols: Optional[int] = None,
+    /,
+    *,
+    k: Optional[int] = 0,
+    dtype: Optional[Dtype] = None,
+    device: Optional[Device] = None,
+) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.eye <numpy.eye>`.
 
@@ -95,15 +146,23 @@ def eye(n_rows: int, n_cols: Optional[int] = None, /, *, k: Optional[int] = 0, d
     from ._array_object import Array
 
     _check_valid_dtype(dtype)
-    if device not in ['cpu', None]:
+    if device not in ["cpu", None]:
         raise ValueError(f"Unsupported device {device!r}")
     return Array._new(np.eye(n_rows, M=n_cols, k=k, dtype=dtype))
 
+
 def from_dlpack(x: object, /) -> Array:
     # Note: dlpack support is not yet implemented on Array
     raise NotImplementedError("DLPack support is not yet implemented")
 
-def full(shape: Union[int, Tuple[int, ...]], fill_value: Union[int, float], *, dtype: Optional[Dtype] = None, device: Optional[Device] = None) -> Array:
+
+def full(
+    shape: Union[int, Tuple[int, ...]],
+    fill_value: Union[int, float],
+    *,
+    dtype: Optional[Dtype] = None,
+    device: Optional[Device] = None,
+) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.full <numpy.full>`.
 
@@ -112,7 +171,7 @@ def full(shape: Union[int, Tuple[int, ...]], fill_value: Union[int, float], *, d
     from ._array_object import Array
 
     _check_valid_dtype(dtype)
-    if device not in ['cpu', None]:
+    if device not in ["cpu", None]:
         raise ValueError(f"Unsupported device {device!r}")
     if isinstance(fill_value, Array) and fill_value.ndim == 0:
         fill_value = fill_value._array
@@ -123,7 +182,15 @@ def full(shape: Union[int, Tuple[int, ...]], fill_value: Union[int, float], *, d
         raise TypeError("Invalid input to full")
     return Array._new(res)
 
-def full_like(x: Array, /, fill_value: Union[int, float], *, dtype: Optional[Dtype] = None, device: Optional[Device] = None) -> Array:
+
+def full_like(
+    x: Array,
+    /,
+    fill_value: Union[int, float],
+    *,
+    dtype: Optional[Dtype] = None,
+    device: Optional[Device] = None,
+) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.full_like <numpy.full_like>`.
 
@@ -132,7 +199,7 @@ def full_like(x: Array, /, fill_value: Union[int, float], *, dtype: Optional[Dty
     from ._array_object import Array
 
     _check_valid_dtype(dtype)
-    if device not in ['cpu', None]:
+    if device not in ["cpu", None]:
         raise ValueError(f"Unsupported device {device!r}")
     res = np.full_like(x._array, fill_value, dtype=dtype)
     if res.dtype not in _all_dtypes:
@@ -141,7 +208,17 @@ def full_like(x: Array, /, fill_value: Union[int, float], *, dtype: Optional[Dty
         raise TypeError("Invalid input to full_like")
     return Array._new(res)
 
-def linspace(start: Union[int, float], stop: Union[int, float], /, num: int, *, dtype: Optional[Dtype] = None, device: Optional[Device] = None, endpoint: bool = True) -> Array:
+
+def linspace(
+    start: Union[int, float],
+    stop: Union[int, float],
+    /,
+    num: int,
+    *,
+    dtype: Optional[Dtype] = None,
+    device: Optional[Device] = None,
+    endpoint: bool = True,
+) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.linspace <numpy.linspace>`.
 
@@ -150,20 +227,31 @@ def linspace(start: Union[int, float], stop: Union[int, float], /, num: int, *,
     from ._array_object import Array
 
     _check_valid_dtype(dtype)
-    if device not in ['cpu', None]:
+    if device not in ["cpu", None]:
         raise ValueError(f"Unsupported device {device!r}")
     return Array._new(np.linspace(start, stop, num, dtype=dtype, endpoint=endpoint))
 
-def meshgrid(*arrays: Sequence[Array], indexing: str = 'xy') -> List[Array, ...]:
+
+def meshgrid(*arrays: Sequence[Array], indexing: str = "xy") -> List[Array, ...]:
     """
     Array API compatible wrapper for :py:func:`np.meshgrid <numpy.meshgrid>`.
 
     See its docstring for more information.
     """
     from ._array_object import Array
-    return [Array._new(array) for array in np.meshgrid(*[a._array for a in arrays], indexing=indexing)]
 
-def ones(shape: Union[int, Tuple[int, ...]], *, dtype: Optional[Dtype] = None, device: Optional[Device] = None) -> Array:
+    return [
+        Array._new(array)
+        for array in np.meshgrid(*[a._array for a in arrays], indexing=indexing)
+    ]
+
+
+def ones(
+    shape: Union[int, Tuple[int, ...]],
+    *,
+    dtype: Optional[Dtype] = None,
+    device: Optional[Device] = None,
+) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.ones <numpy.ones>`.
 
@@ -172,11 +260,14 @@ def ones(shape: Union[int, Tuple[int, ...]], *, dtype: Optional[Dtype] = None, d
     from ._array_object import Array
 
     _check_valid_dtype(dtype)
-    if device not in ['cpu', None]:
+    if device not in ["cpu", None]:
         raise ValueError(f"Unsupported device {device!r}")
     return Array._new(np.ones(shape, dtype=dtype))
 
-def ones_like(x: Array, /, *, dtype: Optional[Dtype] = None, device: Optional[Device] = None) -> Array:
+
+def ones_like(
+    x: Array, /, *, dtype: Optional[Dtype] = None, device: Optional[Device] = None
+) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.ones_like <numpy.ones_like>`.
 
@@ -185,11 +276,17 @@ def ones_like(x: Array, /, *, dtype: Optional[Dtype] = None, device: Optional[De
     from ._array_object import Array
 
     _check_valid_dtype(dtype)
-    if device not in ['cpu', None]:
+    if device not in ["cpu", None]:
         raise ValueError(f"Unsupported device {device!r}")
     return Array._new(np.ones_like(x._array, dtype=dtype))
 
-def zeros(shape: Union[int, Tuple[int, ...]], *, dtype: Optional[Dtype] = None, device: Optional[Device] = None) -> Array:
+
+def zeros(
+    shape: Union[int, Tuple[int, ...]],
+    *,
+    dtype: Optional[Dtype] = None,
+    device: Optional[Device] = None,
+) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.zeros <numpy.zeros>`.
 
@@ -198,11 +295,14 @@ def zeros(shape: Union[int, Tuple[int, ...]], *, dtype: Optional[Dtype] = None,
     from ._array_object import Array
 
     _check_valid_dtype(dtype)
-    if device not in ['cpu', None]:
+    if device not in ["cpu", None]:
         raise ValueError(f"Unsupported device {device!r}")
     return Array._new(np.zeros(shape, dtype=dtype))
 
-def zeros_like(x: Array, /, *, dtype: Optional[Dtype] = None, device: Optional[Device] = None) -> Array:
+
+def zeros_like(
+    x: Array, /, *, dtype: Optional[Dtype] = None, device: Optional[Device] = None
+) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.zeros_like <numpy.zeros_like>`.
 
@@ -211,6 +311,6 @@ def zeros_like(x: Array, /, *, dtype: Optional[Dtype] = None, device: Optional[D
     from ._array_object import Array
 
     _check_valid_dtype(dtype)
-    if device not in ['cpu', None]:
+    if device not in ["cpu", None]:
         raise ValueError(f"Unsupported device {device!r}")
     return Array._new(np.zeros_like(x._array, dtype=dtype))