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.