mcpField
J
jax-development

jax-development

github.com/tristanmanchester開発461 回インストール4 日前に更新
リポジトリ

jax-development について

変換、カーネル、性能を扱うエンジニア向けの JAX 開発スキルです。数学的な前提を明示し、最小の再現テストと対象環境での測定を重視します。

主な機能

  1. 01JAX 環境とプロジェクトの調査
  2. 02最小限の障害再現ケースの作成
  3. 03数値と勾配の正しさの確認
  4. 04コンパイルと再コンパイルの調査
  5. 05入力を制御したベンチマーク
  6. 06シャーディングとマルチホスト動作の確認

活用シーン

  1. 01変換やカーネルの回帰の診断
  2. 02数値および勾配の不変条件の検証
  3. 03コンパイルと実行性能の測定
  4. 04シャーディングと集団通信の評価

インストール

コマンドをコピーしてターミナルで実行します。

$npx skills add https://github.com/tristanmanchester/agent-skills --skill jax-development

インストール後、Claude Code で/jax-developmentを入力して読み込まれたか確認します。

同梱ファイル

SKILL.md と一緒にインストールされる参考資料やスクリプト。

assets/
benchmark_template.py
checkify_template.py
custom_vjp_template.py
export_template.py
issue_report_template.md
mre_template.py
pallas_kernel_skeleton.py
profile_template.py
scan_template.py
shard_map_template.py
sharding_template.py
training_step_template.py
evals/
evals.json
trigger_queries_train.json
trigger_queries_validation.json
evals/files/
bad_control_flow.py
closure_recompile.py
dynamic_mask.py
host_roundtrip.py
legacy_pmap.py
naive_benchmark.py
prng_reuse.py
python_loop_blowup.py
references/
ADVANCED-EXTENSIONS.md
CODE-REVIEW-RUBRIC.md
DEBUGGING-TRIAGE.md
EXPERT-WORKFLOW.md
MENTAL-MODEL.md
PERFORMANCE-PLAYBOOK.md
PORTING-PATTERNS.md

詳細

JAX 開発

変換やカーネルを変更する前に、数学的な契約を明確にします。 形状、dtype、数値許容誤差、乱数の扱い、微分可能性、必須のバックエンドを 記録してください。既存の有用なアーキテクチャは維持してください。新しい API があることだけを 理由に、書き直す必要はありません。互換性のためのシムではなく、現在の公開インターフェースを 使用してください。

調査、再現、計測

SKILL_DIR をこのファイルが含まれるディレクトリに解決してください。同梱スクリプトは 対象プロジェクトの scripts/ ではなく、そのディレクトリにあります。プロジェクトの実際の Python 環境で実行してください。使用する前に、各スクリプトのヘルプを確認してください。

python "$SKILL_DIR/scripts/jax_env_report.py" --format json
python "$SKILL_DIR/scripts/jax_project_scan.py" /absolute/project --format json
python "$SKILL_DIR/scripts/jax_compile_probe.py" --help
python "$SKILL_DIR/scripts/jax_recompile_explorer.py" --help

環境の調査ではバックエンドが初期化され、リソースが割り当てられることがあります。 レポートには関連する環境変数名が記載されますが、その値は記載されません。また、インポート、 バックエンド、スモークテストに失敗した場合は、ゼロ以外の値を返します。これは網羅的な シークレットスキャナーではありません。共有する前に、パス、デバイス、非公開のプロジェクト出力を 確認してください。 静的スキャンで得られるのは手掛かりであり、変換が誤っていることの証明ではありません。モジュールの インポート、トレース、低レベルコードへの変換、ベンチマークではコードが実行されるため、 信頼できる入力を使用してください。

失敗を、関連する形状、dtype、静的引数、変換の順序、バックエンドを維持した最小限の再現例に 絞り込んでください。単純な参照実装と比較し、速度を計測する前に数値的な不変条件と 勾配の不変条件を確認してください。仮説は一度に1つだけ変更し、観測結果を記録してください。

関連するリファレンスを選ぶ

ソースの基準情報には、これらの詳細なリファレンスに影響する現在の変更が記録されています。 オンラインの latest／未リリースのドキュメントを、ロックされたプロジェクトの API とみなさず、 インストール済みのリリースを確認してください。有用なテンプレートや再現・評価用アセットは 同梱されたままです。選択した例を対象環境で実行してください。

正確性とパフォーマンスのルール

通常の変換対象関数は純粋に保ちます。明示的な jax.ref 状態は、独自の変換／副作用に関する制約を持つ、別個のサポート対象モデルであり、隠れた Python のミューテーションを許可するものではありません。型付き PRNG キーを受け渡し、再利用を避けてください。トレースされる判断には構造化制御フローを使用します。小規模な場合は静的な Python ループが適切なこともあるため、すべてのループを機械的に scan に置き換えないでください。

ホストへの転送と同期は意図的に行ってください。入力の準備、トレース／コンパイル、ディスパッチ、デバイス実行、出力の実体化、通信を分けます。x64 と行列積の精度を明示的に確認してください。現行リリースでは、jax.numpy.empty はゼロ初期化されたストレージを保証しなくなりました。ゼロが必要な場合は zeros を使用してください。

ドネーションすると入力バッファーが消費されます。独立した計測試行では、呼び出しごとに新しい入力を構築し、ドネーションしたウォームアップ用引数を再利用しないでください。ベンチマークハーネスでは入力ファクトリーが必須になり、同期エラーが伝播するようになりました。

python "$SKILL_DIR/scripts/jax_benchmark_harness.py" \
  --file /absolute/project/benchmark_case.py --function step --factory make_inputs \
  --jit --donate-argnums 0 --repeat 20

make_inputs() は、同等で新しい入力を含む (args, kwargs) を返します。準備とその同期はタイマーの対象外で、出力の同期は対象内です。形状／dtype／シャーディング／静的値を固定し、出力は別途検証してください。初回呼び出しの時間にはコンパイルと実行が含まれ、キャッシュが使われる場合もあります。純粋なコンパイル時間ではありません。状態を持つ学習のスループットについては、返された状態を引き継ぐ別のベンチマークを作成し、異なるワークロードであることを報告してください。以前の JSON／arrayify や、安全でないドネーション互換経路は残っていません。

目的を達成できる場合は、手動の shard_map よりも、意図的なシャーディングを行うグローバルビューのコードを優先してください。NamedSharding による配置は、明示的な sharding-in-types と同義ではありません。現在のメッシュコンテキストでは with mesh ではなく jax.set_mesh(mesh) を使用します。グローバル／ローカルの形状、レプリケーション、コレクティブ、勾配、出力のシャーディングをテストしてください。1 デバイスでの実行では、マルチホスト通信を検証できません。

根拠を示して仕上げる

診断、パッチ／例、正確性チェック、ハードウェアとワークロードを含む実測時間、残るバックエンドの制限を提示してください。ヘルパーが部分的な出力やエラー出力を返した場合に、ベンチマークやコンパイルが成功したと主張しないでください。ヘルパーのリグレッションについては、python -m unittest discover -s "$SKILL_DIR/tests" -v を実行してください。

コメント

Markdown 対応