ml_dtypes is a stand-alone implementation of several NumPy dtype extensions used in machine learning libraries, including:
bfloat16: an alternative to the standard float16 formatfloat8_e3m4float8_e4m3float8_e4m3b11fnuzfloat8_e4m3fnfloat8_e4m3fnuzfloat8_e5m2float8_e5m2fnuzfloat8_e5m3fnufloat8_e8m0fnufloat4_e2m1fnfloat6_e2m3fnfloat6_e3m2fnint1int2int4uint1uint2uint4complex32 (matching numpy.float16)bcomplex32 (matching numpy.bfloat16)See below for specifications of these number formats.
The ml_dtypes package is tested with Python versions 3.9-3.12, and can be installed with the following command:
pip install ml_dtypes
To test your installation, you can run the following:
pip install absl-py pytest pytest --pyargs ml_dtypes
To build from source, clone the repository and run:
git submodule init git submodule update pip install .
>>> from ml_dtypes import bfloat16 >>> import numpy as np >>> np.zeros(4, dtype=bfloat16) array([0, 0, 0, 0], dtype=bfloat16)
Importing ml_dtypes also registers the data types with numpy, so that they may be referred to by their string name:
>>> np.dtype('bfloat16') dtype(bfloat16) >>> np.dtype('float8_e5m2') dtype(float8_e5m2)
bfloat16A bfloat16 number is a single-precision float truncated at 16 bits.
Exponent: 8, Mantissa: 7, exponent bias: 127. IEEE 754, with NaN and inf.
float4_e2m1fnExponent: 2, Mantissa: 1, bias: 1.
Extended range: no inf, no NaN.
Microscaling format, 4 bits (encoding: 0bSEEM) using byte storage (higher 4 bits are unused). NaN representation is undefined.
Possible absolute values: [0, 0.5, 1, 1.5, 2, 3, 4, 6]
float6_e2m3fnExponent: 2, Mantissa: 3, bias: 1.
Extended range: no inf, no NaN.
Microscaling format, 6 bits (encoding: 0bSEEMMM) using byte storage (higher 2 bits are unused). NaN representation is undefined.
Possible values range: [-7.5; 7.5]
float6_e3m2fnExponent: 3, Mantissa: 2, bias: 3.
Extended range: no inf, no NaN.
Microscaling format, 4 bits (encoding: 0bSEEEMM) using byte storage (higher 2 bits are unused). NaN representation is undefined.
Possible values range: [-28; 28]
float8_e3m4Exponent: 3, Mantissa: 4, bias: 3. IEEE 754, with NaN and inf.
float8_e4m3Exponent: 4, Mantissa: 3, bias: 7. IEEE 754, with NaN and inf.
float8_e4m3b11fnuzExponent: 4, Mantissa: 3, bias: 11.
Extended range: no inf, NaN represented by 0b1000'0000.
float8_e4m3fnExponent: 4, Mantissa: 3, bias: 7.
Extended range: no inf, NaN represented by 0bS111'1111.
The fn suffix is for consistency with the corresponding LLVM/MLIR type, signaling this type is not consistent with IEEE-754. The f indicates it is finite values only. The n indicates it includes NaNs, but only at the outer range.
float8_e4m3fnuz8-bit floating point with 3 bit mantissa.
An 8-bit floating point type with 1 sign bit, 4 bits exponent and 3 bits mantissa. The suffix fnuz is consistent with LLVM/MLIR naming and is derived from the differences to IEEE floating point conventions. F is for “finite” (no infinities), N for with special NaN encoding, UZ for unsigned zero.
This type has the following characteristics:
0bSEEEEMMM0b10000000float8_e5m2Exponent: 5, Mantissa: 2, bias: 15. IEEE 754, with NaN and inf.
float8_e5m2fnuz8-bit floating point with 2 bit mantissa.
An 8-bit floating point type with 1 sign bit, 5 bits exponent and 2 bits mantissa. The suffix fnuz is consistent with LLVM/MLIR naming and is derived from the differences to IEEE floating point conventions. F is for “finite” (no infinities), N for with special NaN encoding, UZ for unsigned zero.
This type has the following characteristics:
0bSEEEEEMM0b10000000float8_e5m3fnuExponent: 5, Mantissa: 3, bias: 15, no sign bit.
Extended range: no inf, NaN represented by 0b11111'111.
float8_e8m0fnuOpenCompute MX scale format E8M0, which has the following properties:
int1, uint1, int2, int4, uint2 and uint41, 2 and 4-bit integer types, where each element is represented unpacked (i.e., padded up to a byte in memory).
NumPy does not support types smaller than a single byte: for example, the distance between adjacent elements in an array (.strides) is expressed as an integer number of bytes. Relaxing this restriction would be a considerable engineering project. These types therefore use an unpacked representation, where each element of the array is padded up to a byte in memory. The lower two or four bits of each byte contain the representation of the number, whereas the remaining upper bits are ignored.
complex32 and bcomplex32Complex types corresponding to numpy.float16 and ml_dtypes.bfloat16 respectively.
Because NumPy currently only recognizes built-in complex types the array attributes arr.imag and arr.real will give incorrect results (.imag returns an array of complex zeros and .real the original values unchanged). To mitigate this issue we provide ml_dtypes.imag() and ml_dtypes.real(). Similarly, other functions may fail to apply the correct complex behavior.
If you're exploring the use of low-precision dtypes in your code, you should be careful to anticipate when the precision loss might lead to surprising results. One example is the behavior of aggregations like sum; consider this bfloat16 summation in NumPy (run with version 1.24.2):
>>> from ml_dtypes import bfloat16 >>> import numpy as np >>> rng = np.random.default_rng(seed=0) >>> vals = rng.uniform(size=10000).astype(bfloat16) >>> vals.sum() 256
The true sum should be close to 5000, but numpy returns exactly 256: this is because bfloat16 does not have the precision to increment 256 by values less than 1:
>>> bfloat16(256) + bfloat16(1) 256
After 256, the next representable value in bfloat16 is 258:
>>> np.nextafter(bfloat16(256), bfloat16(np.inf)) 258
For better results you can specify that the accumulation should happen in a higher-precision type like float32:
>>> vals.sum(dtype='float32').astype(bfloat16) 4992
In contrast to NumPy, projects like JAX which support low-precision arithmetic more natively will often do these kinds of higher-precision accumulations automatically:
>>> import jax.numpy as jnp >>> jnp.array(vals).sum() Array(4992, dtype=bfloat16)
Following NumPy 2's NEP 50 type promotion rules, operations mixing an ml_dtypes array or scalar with a weakly typed Python int, float, or complex scalar convert the Python scalar to the ml_dtypes dtype rather than promoting the operation to float64:
>>> bfloat16(2.0) + 0.0078125 # evaluated in bfloat16, not float64 2 >>> (np.array([1.0, 2.0], dtype=bfloat16) * 0.9).dtype dtype(bfloat16)
Take particular care when comparing Python scalars against non-IEEE formats that lack infinities or signed values. For IEEE formats like bfloat16 or numpy.float16, an out-of-range finite Python float overflows to ±inf, which still preserves ordering against finite values. For finite-only or unsigned formats, however, casting an out-of-range Python scalar either produces NaN (on *fn, *fnuz, and *fnu FP8 types) or saturates to the maximum finite value (on float4_e2m1fn, float6_e2m3fn, and float6_e3m2fn):
>>> from ml_dtypes import float4_e2m1fn, float8_e4m3fn >>> float8_e4m3fn(1000.0) nan >>> float8_e4m3fn(448.0) < 1000.0 # compares 448.0 < NaN False >>> float4_e2m1fn(6.0) == 100.0 # 100.0 saturates to 6.0 True
When comparing values (such as ml_dtypes.finfo(dtype).max) against bounds that may exceed the dtype's representable range, explicitly cast to float or a wider NumPy dtype first:
>>> float(float8_e4m3fn(448.0)) < 1000.0 True
This is not an officially supported Google product.
The ml_dtypes source code is licensed under the Apache 2.0 license (see LICENSE). Pre-compiled wheels are built with the EIGEN project, which is released under the MPL 2.0 license (see LICENSE.eigen).