Skip to content
22 changes: 18 additions & 4 deletions src/array_api_stubs/_draft/data_type_functions.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ def astype(
x: array, dtype: dtype, /, *, copy: bool = True, device: Optional[device] = None
) -> array:
"""
Copies an array to a specified data type irrespective of :ref:`type-promotion` rules.
Copies an array to a specified data type irrespective of :ref:`type-promotion` rules, or to a *kind* of data type.

.. note::
Casting floating-point ``NaN`` and ``infinity`` values to integral data types is not specified and is implementation-dependent.
Expand All @@ -40,8 +40,14 @@ def astype(
----------
x: array
array to cast.
dtype: dtype
desired data type.
dtype_or_kind: Union[dtype, str]
desired data type or kind of data type. Supported kinds are:
- ``'bool'``: boolean data types (e.g., ``bool``).
- ``'signed integer'``: signed integer data types (e.g., ``int8``, ``int16``, ``int32``, ``int64``).
- ``'unsigned integer'``: unsigned integer data types (e.g., ``uint8``, ``uint16``, ``uint32``, ``uint64``).
- ``'integral'``: integer data types. Shorthand for ``('signed integer', 'unsigned integer')``.
- ``'real floating'``: real-valued floating-point data types (e.g., ``float32``, ``float64``).
- ``'complex floating'``: complex floating-point data types (e.g., ``complex64``, ``complex128``).
copy: bool
specifies whether to copy an array when the specified ``dtype`` matches the data type of the input array ``x``. If ``True``, a newly allocated array must always be returned. If ``False`` and the specified ``dtype`` matches the data type of the input array, the input array must be returned; otherwise, a newly allocated array must be returned. Default: ``True``.
device: Optional[device]
Expand All @@ -50,7 +56,15 @@ def astype(
Returns
-------
out: array
an array having the specified data type. The returned array must have the same shape as ``x``.
For ``dtype_or_kind`` a data type, an array having the specified data type.
For ``dtype_or_kind`` a kind of data type:
- If ``x.dtype`` is already of that kind, the data type is maintained.
- Otherwise, an attempt is made to convert to the specified kind, according to the type promotion rules (see :ref:`type-promotion`).
- Numeric kinds are interpreted as the lowest-precision standard data type of that kind for the purposes of type promotion.
For example, ``astype(x, 'complex floating')`` will return an array with the data type ``complex64`` when ``x.dtype`` is ``float32``,
since ``complex64`` is the result of promoting ``float32`` with the lowest-precision standard complex data type, ``complex64``.
- For kind ``integral``, the 'lowest-precision standard data type' is interpreted as ``int8``, not ``uint8``.
The returned array must have the same shape as ``x``.

Notes
-----
Expand Down
Loading