jax.dtypes.canonicalize_dtype

jax.dtypes.canonicalize_dtype(dtype, allow_extended_dtype=False)

Convert from a dtype to a canonical dtype based on config.x64_enabled.

Parameters:

• dtype (Any)

• allow_extended_dtype (bool)

Return type:

dtype | ExtendedDType