このプロジェクトについて

einopsはオープンソースのPythonライブラリで、アインシュタインに着想を得た表記法を通じて柔軟かつ強力なテンソル演算を提供し、テンソル操作のコードをより読みやすく信頼性の高いものにすることを目的としています。この名前は「Einstein-Inspired Notation for operations」に由来し、大まかにnumpy.einsumから着想を得ています。 コアAPIは最小限でありながら強力で、3つの主要な関数を中心に構成されています。rearrange(パターンに従って要素を並べ替える)、reduce(並べ替えと平均や最大などのリダクションを組み合わせる)、repeat(新しい軸に沿ってコピーする)です。その後、異なる次元を持つ複数のテンソルを可逆的に結合するためのpackとunpack、複数文字の軸名をサポートするeinsumも追加されました。einopsには、MLP Mixerや類似のアーキテクチャに適した汎用線形層であるEinMixレイヤーも同梱されており、さらにPyTorch、TensorFlow、Flax、Paddle向けのフレームワーク固有のレイヤーラッパー(Rearrange、Reduceなど)も提供されています。 この表記法は意味情報を重視しています。ユーザーはx.view(x.shape[0], -1)と書く代わりに、rearrange(x, 'b c h w -> b (c h w)')と書くことで、入力と出力の両方の形状を文書化できます。'time c h w -> time (c h w)'のような名前付き軸は、データが独立したバッチではなくビデオシーケンスであることを示唆します。特定の次元(例:c=256、h=19、w=19)を固定することもでき、その場合、操作はコメントに頼るのではなく、実行時に形状を検証します。結果は厳密に決定されるため、depth-to-spaceのような曖昧な操作も明確に記述できます。同じ構文が1D、2D、3Dのプーリングで一貫して機能し、ネイティブAPIが異なる場合でもフレームワーク間で動作は一貫しています(例えば、flattenのセマンティクスはNumPy/PyTorchとTensorFlow/Kerasで異なります)。 einopsは、NumPy、PyTorch、TensorFlow、JAX、CuPy、Flax、Paddle、OneFlow、tinygrad、PyTensorをサポートし、Python配列API標準を実装するあらゆるフレームワーク(NumPy >= 2.0、MLX、pydata/sparse、cubed、nndonnxなど)でも動作します。PyTorchレイヤーについてはtorch.compileとtorch.jit.scriptをサポートしています。このライブラリはPython 3.10以降を必要とし、pip install einopsでインストールできます。このプロジェクトはICLR 2022でオーラル発表され、10,000以上のGitHubプロジェクトで使用されています。