jax-development

Imported from tristanmanchester/agent-skills (jax-development/SKILL.md). Install upstream with npx skills add tristanmanchester/agent-skills --skill jax-development. Copyright stays with the author.

JAX development

Make the mathematical contract explicit before changing transformations or kernels. Record shapes, dtypes, numerical tolerances, randomness, differentiability, and required backend. Preserve useful existing architecture; a new API does not itself justify a rewrite. Use current public interfaces rather than compatibility shims.

Inspect, reproduce, measure

Resolve SKILL_DIR to the directory containing this file. Bundled scripts live there, not in the target project's scripts/. Run them in the project's actual Python environment. Inspect each script's help before use.

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

Environment probing initialises the backend and can allocate resources. The report lists relevant environment variable names, never their values, and returns nonzero for import/backend/smoke-test failure. It is not an exhaustive secret scanner: review paths, devices, and private project output before sharing. Static scans produce leads, not proof that a transformation is wrong. Importing, tracing, lowering, and benchmarking a module execute code; use trusted inputs.

Reduce failures to the smallest reproducer retaining the relevant shape, dtype, static argument, transform order, and backend. Compare against a simple reference and check numerical/gradient invariants before measuring speed. Change one hypothesis at a time and record the observed result.

Select the relevant reference

• Tracing, shapes, pytrees: mental model, transform decisions, and porting patterns.
• Correctness: debugging and review rubric.
• Timing, compilation, memory: performance.
• Placement, collectives, multi-host: sharding.
• Custom derivatives, export, Pallas, FFI: extensions.
• Source-level diagnosis: repo map, workflow, and source baseline.

The source baseline records current changes that affect these longer references. Check the installed release rather than treating online latest/unreleased docs as the API of a locked project. Useful templates and reproducer/evaluation assets remain bundled; run selected examples against the target environment.

Correctness and performance rules

Keep ordinary transformed functions pure; explicit jax.ref state is a separate supported model with its own transform/effect restrictions, not permission for hidden Python mutation. Thread typed PRNG keys and avoid reuse. Use structured control flow for traced decisions; static Python loops can be appropriate when small, so do not mechanically replace every loop with scan.

Keep host transfers and synchronisation deliberate. Separate input preparation, trace/compile, dispatch, device execution, output materialisation, and communication. Check x64 and matmul precision explicitly. jax.numpy.empty no longer promises zero-initialised storage on the current release; use zeros when zeros are needed.

Donation consumes input buffers. For independent timing trials, construct fresh inputs for every call; never reuse a donated warm-up argument. The benchmark harness now requires an input factory and propagates synchronisation errors:

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() returns (args, kwargs) with equivalent, fresh inputs. Preparation and its synchronisation are outside the timer; output synchronisation is inside. Keep shapes/dtypes/shardings/static values fixed and verify outputs separately. First-call time includes compilation/execution and may use caches; it is not pure compile time. For stateful training throughput, write a separate benchmark carrying returned state forward, and report that different workload. No old JSON/arrayify or unsafe donation compatibility path remains.

Prefer global-view code with deliberate sharding before manual shard_map when it meets the objective. NamedSharding placement is not synonymous with explicit sharding-in-types. Current mesh context uses jax.set_mesh(mesh), not with mesh. Test global/local shapes, replication, collectives, gradients, and output sharding; a one-device run cannot validate multi-host communication.

Finish with evidence

Return the diagnosis, patch/example, correctness checks, measured timings with hardware and workload, and remaining backend limitations. Do not claim benchmarks or compilation succeeded when a helper returned partial/error output. For helper regressions run python -m unittest discover -s "$SKILL_DIR/tests" -v.