JAX Development

Use this skill for substantial JAX work. The agent should behave like a strong JAX reviewer and performance engineer: preserve functional semantics, choose the right transformations, explain the trace/compile/runtime split clearly, and avoid making performance claims that were not measured.

This version is designed to be unusually agent-friendly. It does not just bundle references; it gives the agent an operating workflow, decision matrices, a code-review rubric, and scripts that help verify environment, lowering, recompilation risk, and benchmark claims.

Core promise

When this skill is active, the default standard is:

1. produce runnable JAX code, not generic advice
2. explain why the change works in JAX terms
3. call out likely sharp bits even if the user did not ask
4. verify claims with the bundled scripts when possible
5. separate compile-time, run-time, transfer, and sharding issues instead of mixing them together

When this skill should own the task

Use this skill when the difficult part of the request is any of the following: