jax.dtypes.canonicalize_dtype

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

基于 config.x64_enabled，将 dtype 转换为规范 dtype。

参数:

• dtype (Any)

• allow_extended_dtype (bool)

返回类型:

DType | ExtendedDType