ENH: new canonicalize DType function?

[NeilGirdhar]

Jax has canonicalize_dtype and PyTorch also has a notion of default types.

Can we provide canonicalize_dtype for all libraries?

Something like:

def canonicalize_dtype(xp: Namespace, dtype: DType | type[complex]) -> DType:
    if is_jax_namespace(xp):
        from jax.dtypes import canonicalize_dtype
        return canonicalize_dtype(dtype)  # Suppresses warning.
    return xp.empty((), dtype=dtype).dtype

[lucascolley]

Here's JAX's implementation: https://github.com/jax-ml/jax/blob/bb96226dd8a3b4b92431edf2ca75bc82741b4be9/jax/_src/dtypes.py#L296-L310

[lucascolley]

do you have some example code motivating this inclusion?

[NeilGirdhar]

I use it here: https://github.com/NeilGirdhar/efax/blob/658bcbfd8a928a1ead7fa785bd4315f7770fc7ac/efax/_src/parameter.py#L51

[lucascolley]

https://github.com/data-apis/array-api/pull/848/files is in the same general area.

[NeilGirdhar]

I'm not sure I understand how that solves the above problem?

[lucascolley]

You're effectively asking for the default float dtype, right? So xp.astype(xp.empty(()), 'real floating').dtype should achieve the same thing as canonicalize_dtype(float)?

[NeilGirdhar]

[utM] In [1]: import jax.numpy as jnp

[utM] In [2]: x = jnp.zeros(())
WARNING:2025-02-28 12:12:44,915:jax._src.xla_bridge:966: An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.

[utM] In [3]: from array_api_compat import *

[utM] In [4]: xp = array_namespace(x)

[utM] In [5]: xp.astype(xp.empty(()), 'real floating').dtype
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[5], line 1
----> 1 xp.astype(xp.empty(()), 'real floating').dtype

File ~/src/tjax/.venv/lib/python3.13/site-packages/jax/_src/numpy/lax_numpy.py:5632, in astype(x, dtype, copy, device)
   5630 if dtype is None:
   5631   dtype = dtypes.canonicalize_dtype(dtypes.float_)
-> 5632 dtypes.check_user_dtype_supported(dtype, "astype")
   5633 if issubdtype(x_arr.dtype, np.complexfloating):
   5634   if dtypes.isdtype(dtype, ("integral", "real floating")):

File ~/src/tjax/.venv/lib/python3.13/site-packages/jax/_src/dtypes.py:902, in check_user_dtype_supported(dtype, fun_name)
    898   warnings.warn("Passing an array as a dtype argument is deprecated; "
    899                 "instead of dtype=arr use dtype=arr.dtype.",
    900                 category=DeprecationWarning, stacklevel=3)
    901   return  # no further check needed, as array dtypes have already been validated.
--> 902 if issubdtype(dtype, extended):
    903   return
    904 # Avoid using `dtype in [...]` because of numpy dtype equality overloading.

File ~/src/tjax/.venv/lib/python3.13/site-packages/jax/_src/dtypes.py:431, in issubdtype(a, b)
    414 """Returns True if first argument is a typecode lower/equal in type hierarchy.
    415 
    416 This is like :func:`numpy.issubdtype`, but can handle dtype extensions such as
    417 :obj:`jax.dtypes.bfloat16` and `jax.dtypes.prng_key`.
    418 """
    419 # Main departures from np.issubdtype are:
    420 # - "extended" dtypes (like prng key types) are not normal numpy dtypes, so we
    421 #   need to handle them specifically. However, their scalar types do conform to
   (...)    428 # unhashable (e.g. custom objects with a dtype attribute). The following check is
    429 # fast and covers the majority of calls to this function within JAX library code.
    430 return _issubdtype_cached(
--> 431   a if isinstance(a, _types_for_issubdtype) else np.dtype(a),  # type: ignore[arg-type]
    432   b if isinstance(b, _types_for_issubdtype) else np.dtype(b),  # type: ignore[arg-type]
    433 )

TypeError: data type 'real floating' not understood

[lucascolley]

TypeError: data type 'real floating' not understood

That's what data-apis/array-api#848 is changing!

[NeilGirdhar]

Oh, my mistake! I guess if you're okay with this, you can just do xp.empty((), dtype=dtype).dtype. I'll close since this is a one-liner, although it's a bit abstruse.