tinygrad/tinygrad

tinygrad – 小さくハック可能なディープラーニングスタック

概要 – tinygradは、完全なエンドツーエンドのディープラーニングフレームワークを実装したオープンソースのPythonライブラリですが、コードベースは意図的に小さく保たれています。以下を提供します：

自動微分 (autograd) を備えた Tensor クラス。
演算を融合し、デバイスカーネルに低レベル化する IR / コンパイラ。
JIT / グラフ実行 (TinyJit 経由)。
実際のモデルを学習するための軽量な nn / optim / dataset 層。

PyTorchのような人間工学的な使い勝手を提供しつつ、コンパイラの内部構造を公開して調整を容易にすることで、フル機能のPyTorchと極小のmicrogradの中間に位置することを目指しています。

主要フレームワークとの比較

機能	PyTorch	JAX	tinygrad
Eager Tensor API	✅	–	✅
Autograd	✅	✅ (via JAXPR)	✅
関数型変換 (vmap, pmap)	–	✅	部分対応
コンパイラとIRの可視性	–	✅	✅
フロントエンド (nn, optim, datasets)	✅	–	✅
マルチデバイスサポート	✅	✅	✅ (CPU, CUDA, Metal, OpenCL, AMD, NV, QCOM, WebGPU)

クイックスタート例

from tinygrad import Tensor, nn, Context

class LinearNet:
    def __init__(self):
        self.l1 = Tensor.kaiming_uniform(784, 128)
        self.l2 = Tensor.kaiming_uniform(128, 10)
    def __call__(self, x: Tensor) -> Tensor:
        return x.flatten(1).dot(self.l1).relu().dot(self.l2)

model = LinearNet()
optim = nn.optim.Adam([model.l1, model.l2], lr=0.001)

x, y = Tensor.rand(4, 1, 28, 28), Tensor([2,4,3,7])
with Context(TRAINING=1):
    for i in range(10):
        optim.zero_grad()
        loss = model(x).sparse_categorical_crossentropy(y).backward()
        optim.step()
        print(i, loss.item())

リポジトリには、GPU上で約5秒で約98%の精度に達する完全なMNISTの例 (examples/beautiful_mnist.py) が含まれています。

デバイスサポート (アクセラレータ)

tinygradは、デバイスあたり約25個のプリミティブ演算のみで多くのバックエンド上で動作します：

CPU
CUDA (NVIDIA GPU)
Metal (Apple GPU)
OpenCL
AMD, NV, QCOM (特殊なランタイム)
WebGPU (ブラウザ / WASM) python3 -c "from tinygrad import Device; print(Device.DEFAULT)" を実行して、デフォルトで使用されるアクセラレータを確認してください。

インストール

# ソースからインストール (推奨)
git clone https://github.com/tinygrad/tinygrad.git
cd tinygrad
python3 -m pip install -e .

# またはリポジトリから直接インストール
python3 -m pip install git+https://github.com/tinygrad/tinygrad.git

プロジェクトは、テストスイートを実行するための .[testing] エクストラも提供しています。

ドキュメントとコミュニティ

ドキュメントとクイックスタート: https://docs.tinygrad.org/ (docs/ フォルダから生成)
Discord: https://discord.gg/ZjZadyC7PK – ユーザーと開発者のための活発なチャット。
貢献: READMEには厳格なPRポリシー（コードゴルフの禁止、速度に関する主張のベンチマーク、回帰テストの追加など）が記載されており、採用された改善には報奨金が支払われることもあります。

tinygradの用途

学習と実験 – コンパイラのパイプラインが公開されているため、高レベルのテンソル演算がどのように低レベルのカーネルになるかを正確に確認できます。
限られたハードウェアでのプロトタイピング – ライブラリは軽量で、CPU、GPU、さらにはWebGPUでも動作します。
変更可能なスタックを必要とする研究 – 数個のファイルを変更するだけで、新しい演算を追加したりスケジューリングを変更したりできます。

巨大なエコシステム、本番環境レベルの安定性、またはJAXの変換機能のフルセットが必要な場合は、PyTorchやJAXの方が適しています。

要約

tinygradは、テンソル、自動微分、JITベースのコンパイラ、マルチデバイス実行をコンパクトで読みやすいコードベースで実装した、実用的で機能的なディープラーニングフレームワークです。巨大なライブラリのオーバーヘッドなしに、現代のMLスタックの内部構造を理解したり拡張したりしたい学生、愛好家、研究者に最適です。