mcpField
J
jax-development

jax-development

github.com/tristanmanchester开发461 次安装4 天前更新
代码仓库

关于 jax-development

面向 JAX 工程开发的技能，适用于修改转换、内核和性能的工程师。它强调明确数学约束、构建最小可复现测试，并在目标项目环境中进行测量。

主要功能

  1. 01检查 JAX 环境与项目
  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 开发

在更改变换或内核之前，先明确数学契约。 记录形状、数据类型、数值容差、随机性、可微性和所需后端。 保留现有架构中有用的部分；仅凭新增 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

环境探测会初始化后端并可能分配资源。报告会列出相关环境变量的名称，但不会列出其值；若导入、后端或冒烟测试失败，则返回非零状态。它不是全面的机密扫描器：分享前请检查路径、设备和项目私有输出。静态扫描只能提供线索，不能证明某个变换有误。导入、追踪、降级和基准测试模块时都会执行代码；请使用可信输入。

将故障缩减为最小复现，同时保留相关形状、数据类型、静态参数、变换顺序和后端。与简单参考实现进行比较，并在测量速度前检查数值和梯度不变量。一次只更改一个假设，并记录观察结果。

选择相关参考资料

源码基线记录了会影响这些较长参考资料的当前变更。请检查已安装的发行版，不要将在线的 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 放置并不等同于类型中的显式分片。当前网格上下文使用 jax.set_mesh(mesh)，而不是 with mesh。测试全局／本地形状、复制、集合通信、梯度和输出分片；单设备运行无法验证多主机通信。

提供证据后结束

给出诊断、补丁／示例、正确性检查、包含硬件和工作负载信息的实测时间，以及剩余的后端限制。如果辅助函数返回了部分结果／错误输出，不要声称基准测试或编译成功。对于辅助函数回归，请运行 python -m unittest discover -s "$SKILL_DIR/tests" -v。

评论

支持 Markdown