Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: astype: accept a kind of data type #848

Open
wants to merge 11 commits into
base: main
Choose a base branch
from
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``).
lucascolley marked this conversation as resolved.
Show resolved Hide resolved
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