jax-ml/ml_dtypes

何を解決するか

NumPy ではネイティブにサポートされていないが、機械学習でよく使われる特殊な数値データ型（dtypes）の独立した実装を提供します。これにより、ML ライブラリは低精度フォーマットを使用してメモリ使用量を削減し、計算効率を向上させることができます。

動作方法

このライブラリは、NumPy の dtype 拡張を複数実装し、NumPy 配列内で直接使用できるように登録します。以下の広範なフォーマットをサポートしています：

• Bfloat16：単精度浮動小数点の短縮版。
• 8ビット浮動小数点：指数部と仮数部のビット数のさまざまな構成（例：float8_e4m3、float8_e5m2）。
• マイクロスケーリング（MX）フォーマット：バイト未満の浮動小数点表現（4ビットおよび6ビット）。
• 狭い整数：1、2、4ビットの整数型（展開されたバイトとして格納）。
• 複素数型：16ビットの複素浮動小数点数（complex32 および bcomplex32）。

対象ユーザー

NumPy のコアエンジンを再実装せずに、学習や推論に低精度数値フォーマットを使用したい機械学習ライブラリ開発者や研究者。

特徴

• NumPy との統合：NumPy に dtype を登録し、文字列名で参照可能にします。
• 広範なフォーマット対応：8ビット、4ビット、6ビットの浮動小数点および整数フォーマットを多数実装しています。
• 低精度処理：算術演算における精度損失を管理するためのガイドラインとツールを提供し、高精度の蓄積を推奨します。