Add meshgrid(), broadcast_arrays(), broadcast_to(), and can_cast() to the array API namespace

4 files changed, 50 insertions, 11 deletions

diff --git a/numpy/_array_api/__init__.py b/numpy/_array_api/__init__.py
index ebbe0bb91..56699b09d 100644
--- a/numpy/_array_api/__init__.py
+++ b/numpy/_array_api/__init__.py
@@ -116,13 +116,13 @@ from ._constants import e, inf, nan, pi
 
 __all__ += ['e', 'inf', 'nan', 'pi']
 
-from ._creation_functions import asarray, arange, empty, empty_like, eye, from_dlpack, full, full_like, linspace, ones, ones_like, zeros, zeros_like
+from ._creation_functions import asarray, arange, empty, empty_like, eye, from_dlpack, full, full_like, linspace, meshgrid, ones, ones_like, zeros, zeros_like
 
-__all__ += ['asarray', 'arange', 'empty', 'empty_like', 'eye', 'from_dlpack', 'full', 'full_like', 'linspace', 'ones', 'ones_like', 'zeros', 'zeros_like']
+__all__ += ['asarray', 'arange', 'empty', 'empty_like', 'eye', 'from_dlpack', 'full', 'full_like', 'linspace', 'meshgrid', 'ones', 'ones_like', 'zeros', 'zeros_like']
 
-from ._data_type_functions import finfo, iinfo, result_type
+from ._data_type_functions import broadcast_arrays, broadcast_to, can_cast, finfo, iinfo, result_type
 
-__all__ += ['finfo', 'iinfo', 'result_type']
+__all__ += ['broadcast_arrays', 'broadcast_to', 'can_cast', 'finfo', 'iinfo', 'result_type']
 
 from ._dtypes import int8, int16, int32, int64, uint8, uint16, uint32, uint64, float32, float64, bool
 
diff --git a/numpy/_array_api/_creation_functions.py b/numpy/_array_api/_creation_functions.py
index 003b10afb..c6db3cb7b 100644
--- a/numpy/_array_api/_creation_functions.py
+++ b/numpy/_array_api/_creation_functions.py
@@ -3,8 +3,10 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 if TYPE_CHECKING:
-    from ._types import (Optional, SupportsDLPack, SupportsBufferProtocol, Tuple,
-                         Union, array, device, dtype)
+    from ._types import (List, Optional, SupportsDLPack,
+                         SupportsBufferProtocol, Tuple, Union, array, device,
+                         dtype)
+    from collections.abc import Sequence
 from ._dtypes import _all_dtypes
 
 import numpy as np
@@ -135,6 +137,15 @@ def linspace(start: Union[int, float], stop: Union[int, float], num: int, /, *,
         raise NotImplementedError("Device support is not yet implemented")
     return ndarray._new(np.linspace(start, stop, num, dtype=dtype, endpoint=endpoint))
 
+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 ndarray
+    return [ndarray._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>`.
diff --git a/numpy/_array_api/_data_type_functions.py b/numpy/_array_api/_data_type_functions.py
index d4816a41f..81eacfe0f 100644
--- a/numpy/_array_api/_data_type_functions.py
+++ b/numpy/_array_api/_data_type_functions.py
@@ -4,12 +4,40 @@ from ._array_object import ndarray
 
 from typing import TYPE_CHECKING
 if TYPE_CHECKING:
-    from ._types import Union, array, dtype
-
-from collections.abc import Sequence
+    from ._types import List, Tuple, Union, array, dtype
+    from collections.abc import Sequence
 
 import numpy as np
 
+def broadcast_arrays(*arrays: Sequence[array]) -> List[array]:
+    """
+    Array API compatible wrapper for :py:func:`np.broadcast_arrays <numpy.broadcast_arrays>`.
+
+    See its docstring for more information.
+    """
+    from ._array_object import ndarray
+    return [ndarray._new(array) for array in np.broadcast_arrays(*[a._array for a in arrays])]
+
+def broadcast_to(x: array, shape: Tuple[int, ...], /) -> array:
+    """
+    Array API compatible wrapper for :py:func:`np.broadcast_to <numpy.broadcast_to>`.
+
+    See its docstring for more information.
+    """
+    from ._array_object import ndarray
+    return ndarray._new(np.broadcast_to(x._array, shape))
+
+def can_cast(from_: Union[dtype, array], to: dtype, /) -> bool:
+    """
+    Array API compatible wrapper for :py:func:`np.can_cast <numpy.can_cast>`.
+
+    See its docstring for more information.
+    """
+    from ._array_object import ndarray
+    if isinstance(from_, ndarray):
+        from_ = from_._array
+    return np.can_cast(from_, to)
+
 def finfo(type: Union[dtype, array], /) -> finfo_object:
     """
     Array API compatible wrapper for :py:func:`np.finfo <numpy.finfo>`.
diff --git a/numpy/_array_api/_types.py b/numpy/_array_api/_types.py
index 32d03e2a7..36c9aa610 100644
--- a/numpy/_array_api/_types.py
+++ b/numpy/_array_api/_types.py
@@ -6,10 +6,10 @@ annotations in the function signatures. The functions in the module are only
 valid for inputs that match the given type annotations.
 """
 
-__all__ = ['Literal', 'Optional', 'Tuple', 'Union', 'array', 'device',
+__all__ = ['List', 'Literal', 'Optional', 'Tuple', 'Union', 'array', 'device',
            'dtype', 'SupportsDLPack', 'SupportsBufferProtocol', 'PyCapsule']
 
-from typing import Literal, Optional, Tuple, Union, TypeVar
+from typing import List, Literal, Optional, Tuple, Union, TypeVar
 
 from . import (ndarray, int8, int16, int32, int64, uint8, uint16, uint32,
                uint64, float32, float64)