ml_dtypes

ml_dtypes is a stand-alone implementation of several NumPy dtype extensions used across machine learning libraries. It adds the low-precision numeric formats that modern ML frameworks rely on — most notably bfloat16, a family of 8-bit floating point types (float8_e4m3fn, float8_e5m2, and more), Microscaling (MX) sub-byte floats, and narrow integer encodings like int4 and uint4.

Each format is registered as a proper NumPy dtype, so arrays of these types behave like native NumPy arrays with correct casting, rounding, and special-value handling. ml_dtypes is a foundational dependency of JAX, TensorFlow, and other frameworks, giving the whole ecosystem a single, consistent implementation of these ML number formats.

What You Get

• A NumPy-registered bfloat16 dtype
• A full family of 8-bit float formats (float8_e4m3fn, float8_e5m2, and more)
• Microscaling (MX) sub-byte float formats (float4/float6 variants)
• Narrow integer encodings (int1, int2, int4 and unsigned variants)
• Correct casting, rounding, and NaN/Inf handling for each format
• A shared, framework-agnostic implementation used by JAX and TensorFlow

Common Use Cases

• Working with bfloat16 or float8 arrays in NumPy for ML experiments
• Prototyping low-precision quantization schemes before deploying to accelerators
• Sharing a consistent dtype implementation across ML frameworks
• Inspecting or converting model weights stored in reduced-precision formats
• Emulating accelerator number formats on CPU with NumPy

Under The Hood

Architecture — ml_dtypes registers new scalar types with NumPy’s type system so that arrays of bfloat16, the various float8_* formats, MX sub-byte floats, and narrow integer encodings behave like first-class dtypes. The numerical cores are implemented in C++ (headers under ml_dtypes/include, e.g. float8.h), templated on exponent/mantissa/bias parameters, and exposed to Python through a compiled extension that wires up casts, ufunc loops, and array-scalar behavior. A thin Python package re-exports the registered dtypes and helper metadata.

Tech Stack — The heavy lifting is C++ built as a Python extension (the repo’s primary language is C++), with a NumPy dependency for the array integration and pybind-style bindings. Builds are driven by a modern pyproject.toml and produce wheels across platforms via GitHub Actions; the library targets Python 3.9–3.12.

Code Quality — The project has a dedicated test workflow and wheel-build CI, with tests validating each numeric format’s rounding, special values (NaN/Inf), and cast behavior. Because it underpins JAX, TensorFlow, and other frameworks, correctness of the low-level float representations is treated rigorously.

API Design — The public surface is intentionally tiny: import ml_dtypes and use ml_dtypes.bfloat16, ml_dtypes.float8_e4m3fn, etc. directly as NumPy dtypes. There is almost no boilerplate — the dtypes drop straight into existing NumPy code — and the README documents each format’s bit layout precisely.