RFC: add data type inspection utilities to the array API specification

[kgryte]

This RFC proposes adding data type inspection utilities to the array API specification.

Overview

Currently, the array API specification requires that conforming implementations provide a specified set of data type objects (see https://data-apis.org/array-api/2021.12/API_specification/data_types.html) and casting functions (see https://data-apis.org/array-api/2021.12/API_specification/data_type_functions.html).

However, the specification does not include APIs for array data type inspection (e.g., an API for determining whether an array has a complex number data type or a floating-point data type, etc).

Prior Art

NumPy and its derivatives have dtype objects with extensive properties, including a kind property, which returns a character code indicating the general "kind" of data. For example, for relevant dtypes in the specification, NumPy uses the following character codes:

• b: boolean
• i: signed integer
• u: unsigned integer
• f: floating-point (real-valued)
• c: complex floating-point

In [1]: np.zeros((3,4)).dtype.kind
Out[1]: 'f'

This availability of the kind property is useful when wanting to branch based on input array data types (e.g., applying summation algorithms).

if x.dtype.kind == 'f':
    # do one thing
else:
   # do another thing

In PyTorch, dtype objects have is_complex and is_floating_point properties for checking a data type "kind".

Additionally, PyTorch offers functional APIs is_complex and is_floating_point providing equivalent behavior.

Proposal

Given the proposal for adding complex number support to the specification (see #373 and #418), a greater need arises for the specification to require conforming implementations to provide standardized ways for data type inspection.

For example, conforming implementations will need to branch in abs(x) depending on whether x is real-valued or complex-valued. Similarly, in downstream user code, we can expect that users will inevitably encounter situations where they need to branch based on input array data types (e.g., when choosing summation algorithms).

As this specification has favored functional APIs, this RFC follows suit and proposes adding the following APIs to the specification:

has_complex_float_dtype(x: Union[array, dtype]) -> bool

Returns a bool indicating whether an input array has a complex number data type (e.g., complex64 or complex128).

has_real_float_dtype(x: Union[array, dtype]) -> bool

Returns a bool indicating whether an input array has a (real-valued) floating-point number data type (e.g., float32 or float64).

has_float_dtype(x: Union[array, dtype]) -> bool

Returns a bool indicating whether an input array has a complex or real-valued floating-point number data type (e.g., float32, float64, complex64, or complex128).

has_unsigned_int_dtype(x: Union[array, dtype]) -> bool

Returns a bool indicating whether an input array has an unsigned integer data type (e.g., uint8, uint16, uint32, uint64).

has_signed_int_dtype(x: Union[array, dtype]) -> bool

Returns a bool indicating whether an input array has a signed integer data type (e.g., int8, int16, int32, int64).

has_int_dtype(x: Union[array, dtype]) -> bool

Returns a bool indicating whether an input array has an integer (signed or unsigned) data type.

has_real_dtype(x: Union[array, dtype]) -> bool

Returns a bool indicating whether an input array has a real-valued (integer or floating-point) data type.

has_bool_dtype(x: Union[array, dtype]) -> bool

Returns a bool indicating whether an input array has a boolean data type.

---

The above APIs cover the list of data types currently described in the specification, are sufficiently specific to cover most use cases, and can be composed to address most anticipated data type set combinations.

[jakirkham]

Also related ( #152 )

Edit: If we look into type naming (briefly discussed), this discussion around typing naming in Zarr may be of interest ( zarr-developers/zarr-specs#131 )

[honno]

Like can_cast() or result_type(), could these utils take both dtypes and arrays? I'd personally want these utils for dtype objects themselves, but definitely my own use cases are not quite aligned with most array consumers.

[kgryte]

Update: I've updated the OP as follows based on feedback here and in the last array API consortium meeting.

1. Functions now accept both arrays and dtypes.
2. Function names include a _dtype suffix (as suggested during the consortium meeting)
3. Function names begin with a has_ prefix. This helps avoid conflicts with existing APIs (e.g., PyTorch) and matches how one might describe an array (e.g., has shape X, has data type Y, etc).
4. Included both real-valued and generic float APIs to match specification data type categories.
5. Included a generic real dtype API to match specification data type categories.

[rgommers]

Some more prior art:

• TensorFlow has methods like is_bool on DType objects: https://www.tensorflow.org/api_docs/python/tf/dtypes/DType
• JAX doesn't have anything other than a subset of numpy APIs (result_type, can_cast, promote_dtypes): https://jax.readthedocs.io/en/latest/_modules/jax/_src/dtypes.html
• NumPy issue on the too many ways of comparing: ENH: add a canonical way to determine if dtype is integer, floating point or complex numpy/numpy#17325

As this specification has favored functional APIs

Given that dtype objects are immutable and have no state, this should also work for JAX et al. Not saying that that's my preference (I'm not yet sure), but this RFC proposes a lot of functions ...

3 underscores in a name like has_real_float_dtype is also not ideal.

[kgryte]

Given that dtype objects are immutable and have no state, this should also work for JAX et al. Not saying that that's my preference (I'm not yet sure), but this RFC proposes a lot of functions ...

Whether methods or functions, surface area would be the same. The list can obviously be culled; however, I do think there is some advantage to matching the categories as used in the spec, especially for providing consistent APIs for input argument validation.

3 underscores in a name like has_real_float_dtype is also not ideal.

The number of underscores is not super important, IMO. Instead, we're probably concerned about number of characters. Originally, I left out the _dtype suffix, which would reduce the function name length; however, consortium members voiced desire for such a suffix in the array API meeting.

I don't have a strong opinion here; although, the current naming convention is arguably more literate.

[vnmabus]

Silly question, why not do:

if array.dtype in <set_of_dtypes>:
    ...

and require implementations to provide some predefined sets, such as "set of all supported integer dtypes" or "set of all supported floating point dtypes"?

[rgommers]

and require implementations to provide some predefined sets, such as "set of all supported integer dtypes" or "set of all supported floating point dtypes"?

That does seem more appealing indeed; it's what can already be done today and it reads fairly well. I think I prefer that over both the has_* functions in this proposal and the numpy issubdtype design.

[seberg]

I don't like issubdtype. For NumPy, I could imagine isinstance(arr.dtype, InexactDType) (or similar). So that way the API here would be isinstance(arr.dtype, some_object). The problem is that I am not sure if an isinstance API would work for everyone.

For arr.dtype in set_of_dtypes there are two three things to keep in mind:

1. The set_of_dtypes will be different for each library, because bfloat16, float16, and others do not exist for implementers. Implementers can extend the API after all.
2. For NumPy, users may extend the API reasonably soon. For example adding bfloat16 or a multi-precision float object.
3. Sets might be tricky right now NumPy in either case (although that could likely be made sure to work). There are arbitrary number of possible instances for dtypes, although they should compare equal with a limited set, that set is confusingly large (byte-order matters).

I do think neither of these is particularly problematic. But, I would say that this would not be a set, but rather an opaque object that supports the in operator.

[honno]

A minor pro of dtype sets is that it could be a way for a library to communicate what dtypes they support—thinking of PyTorch and it only supporting uint8 unsigned integers. Useful here and there, like telling Hypothesis to not try generating uint{16/32/64}.

[rgommers]

@leofang point out that this is blocking for adding real and conj (and I imagine imag too), it'd be great to finalize this. The majority of folks who have weighed in seem to prefer a set/collection type of approach. So here's a suggested API for that, in line with @seberg's last sentence above.

1. There must be objects integer_dtypes, floating_point_dtypes, and complex_dtypes,
2. The syntax dtype in xxx_dtypes must yield a boolean value with the expected result (to be detailed out more in the spec),
3. xxx_dtypes must contain all the expected dtypes that are part of the standard, and may contain additional dtypes of the same kind
4. The objects may be of any kind, e.g. a set or a custom class instance.

Other thoughts:

• No object for boolean is needed, because bool already supports __eq__, so array.dtype == bool is enough.
• Also no separate signed/unsigned integer objects, because that's a bit much for the API / less needed. This is mostly a convenient way to spell array.dtype in (dtype1, dtype2, ...) anyway.
• The one name where there's not a single obvious choice is floating_point_dtypes. It could also be float_dtypes, floating_dtypes, or real_dtypes for example.
• Not specifying the type of these objects is on purpose, to make it easy to for example have an API that adds user-defined dtypes in.
  • That means that for static typing we need another Protocol. Not completely ideal, but imho better than restricting implementation choices for libraries (see point 3 in @seberg's comment above about why set is tricky for NumPy).

One alternative with a similar API surface is to add 3 functions with the same functionality instead. Those functions could be 3 of the ones in the issue description here (e.g., has_integer_dtype, has_floating_point_dtype, has_complex_dtype). Considerations:

• Pro: it's better for static typing,
• Con: it introduces an asymmetry between supported and unsupported sets - we need the dtype in xxx anyway when the predefined objects aren't the right ones.

I think the con is more important than the pro here. But I'd say either choice is pretty reasonable here.

[rgommers]

Just to make sure @leofang, both flavors are fine for accelerators, right? When the spec says something should return a bool, that's not a problem - only Python control flow like if _expr_yielding_a_bool is. So a function is not preferred from that perspective. Or maybe there's a significant amount of extra implementation complexity for the dtype in xxx_dtypes version?