data-apis/array-api-extra

ENH: real and complex dtype functions

izaid opened this issue · 3 comments

Hello! Very excited about this library, as it (hopefully) provides a place to put functions that are array agnostic and usable by many people. Let me describe two simple ones that I'd like to make a PR for.

I'd like to add functions real_type(dtype) and complex_type(dtype). These are simple:

  • real_type takes in a real or complex dtype, and returns the corresponding floating-point dtype. So complex128 -> float64, complex64 -> float32, float64 -> float64, float32 -> float32.
  • complex_type takes in a real or complex dtype, and returns the corresponding complex floating-point dtype. So float64 -> complex128, float32 -> complex64, complex128 -> complex128, complex64 -> complex64.

This can be really helpful in application code to work out the correct types. Shall I make a PR for these?

consumer libraries definitely want things along these lines, yes. I recently wrote this in SciPy:

def xp_float_to_complex(arr: Array, xp: ModuleType | None = None) -> Array:
    xp = array_namespace(arr) if xp is None else xp
    arr_dtype = arr.dtype
    # The standard float dtypes are float32 and float64.
    # Convert float32 to complex64,
    # and float64 (and non-standard real dtypes) to complex128
    if xp.isdtype(arr_dtype, xp.float32):
        arr = xp.astype(arr, xp.complex64)
    elif xp.isdtype(arr_dtype, 'real floating'):
        arr = xp.astype(arr, xp.complex128)

    return arr

and there is also the more elaborate xp_broadcast_promote in https://github.com/scipy/scipy/blob/main/scipy/_lib/_array_api.py.

Could either of those functions be written more simply using the functions you propose?

I would probably suggest hashing out the behaviour more fully before opening a PR. What should the functions do for other standard dtypes like int64? Also, I think it would be in the spirit of this library (in its current form) to error on any object that isn’t a standard dtype.

There was also some recent discussion about what a standardized API for this could look like data-apis/array-api#841 (comment)

@izaid would these functions still be useful to you after data-apis/array-api#848? That covers the case of casting from potentially real or complex input to the appropriate complex data type via xp.astype(x, 'complex floating').