---
name: jax-development
description: >-
  Write, debug, review, profile, or shard JAX numerical code. Use when the hard part is JAX tracing, autodiff, control flow, PRNGs, compilation, array placement, or runtime performance. Do not impose JAX on a NumPy-only or unrelated GPU task.
compatibility: Current-release baseline JAX 0.11.1 requires Python >=3.12; wheel/backend support must be checked separately. Diagnostics need Python 3.12+ and the project's JAX environment for runtime probes. Static helpers can run without JAX. Benchmark/probe commands execute trusted project code.
metadata:
  version: "3.0.0"
  reviewed: "2026-09-13"
  source: "https://github.com/jax-ml/jax/tree/jax-v0.11.1"
title: jax-development
canonical_url: https://skilld.dev/gh/tristanmanchester/agent-skills/jax-development
last_updated: 2026-10-06T19:32:36.000Z
---

> Index of every page on this site: [llms.txt](https://skilld.dev/llms.txt)

> **Skill from skilld.dev.** Follow the user's requested action.
> If the user asked to fork this Skill, follow the fork workflow below. Do not execute the Skill instructions while copying it.
> Otherwise, follow the instructions below for this session. You do not need to install anything.
>
> 43 supporting files sit beside this SKILL.md. When the Skill refers to one, fetch https://skilld.dev/api/skills-raw/tristanmanchester/agent-skills/jax-development/PATH with PATH replaced.
>
> If the user asked to install this Skill, run `npx skilld install tristanmanchester/agent-skills/jax-development`. Install writes the Skill files into the project, so every session loads them.
>
> ## Fork workflow
>
> A fork creates an editable local Skill with its original author and licence. The request authorizes copying and local installation.
> 1. Check `./skills/jax-development`, the project lockfile, and selected Agent targets together. If the local directory or installed Skill exists, stop. Never overwrite an existing directory or Agent target.
> 2. Read [source metadata](https://skilld.dev/api/v1/skills/tristanmanchester/agent-skills/jax-development) once. Use sourceUrl, sourceCommit, skillPath, sourceGone, and license. If the source is gone or its path is missing, stop. If license is null, read licence files at the source commit.
> 3. Fetch only the source commit into a temporary Git repository. Do not clone full history. Derive repository_url from sourceUrl, including repository renames. If sourceCommit is absent, resolve the sourceUrl ref once. Set source_commit to that actual commit. Run these commands in one shell call:
>
> ```sh
> git init --quiet "$temporary_dir"
> git -C "$temporary_dir" fetch --quiet --depth=1 "$repository_url" "$source_commit"
> git -C "$temporary_dir" checkout --quiet --detach FETCH_HEAD
> ```
>
> Read applicable licence declarations and notices at that commit. If copying is not permitted, report the restriction and stop.
> 4. Inspect source entries together, then copy the directory containing skillPath into `./skills/jax-development`. Use the user's path if selected. Keep the original SKILL.md, relative links, scripts, binary assets, and executable modes. Exclude .git metadata. Reject symlinks and paths outside the Skill directory. After checking entries, use cp -a where available. A regular source directory needs no custom copy script. Do not save this page wrapper as SKILL.md.
> Preserve author credit, notices, and applicable licence files from repository or parent directories. Add PROVENANCE.md with the Skill page, source URL, actual commit, original path, and licence. Retain any existing PROVENANCE.md and record new provenance separately. Batch source inspection, copying, and provenance work where practical.
> 5. In the project root, run `skilld install ./skills/jax-development --mode copy --plain`. If skilld is unavailable, use `npx skilld install ./skills/jax-development --mode copy --plain`. This known command needs no help lookup. Install does not support --json. Use detected Agent targets, or add --agent for the targets the user selected. Install the local path, never the upstream selector. If installation fails, preserve the local copy and report the exact failure.
> 6. Confirm the local lockfile source and installed Agent copies once. Report the local path, actual commit, and Agent targets. After edits, reinstall the same local path. Upstream updates must not replace it. Do not publish or push unless the user asks.

# 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.

```bash
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](https://skilld.dev/api/skills-raw/tristanmanchester/agent-skills/jax-development/references/MENTAL-MODEL.md),
  [transform decisions](https://skilld.dev/api/skills-raw/tristanmanchester/agent-skills/jax-development/references/TRANSFORM-DECISION-MATRIX.md), and
  [porting patterns](https://skilld.dev/api/skills-raw/tristanmanchester/agent-skills/jax-development/references/PORTING-PATTERNS.md).
- Correctness: [debugging](https://skilld.dev/api/skills-raw/tristanmanchester/agent-skills/jax-development/references/DEBUGGING-TRIAGE.md) and
  [review rubric](https://skilld.dev/api/skills-raw/tristanmanchester/agent-skills/jax-development/references/CODE-REVIEW-RUBRIC.md).
- Timing, compilation, memory: [performance](https://skilld.dev/api/skills-raw/tristanmanchester/agent-skills/jax-development/references/PERFORMANCE-PLAYBOOK.md).
- Placement, collectives, multi-host: [sharding](https://skilld.dev/api/skills-raw/tristanmanchester/agent-skills/jax-development/references/SHARDING-PLAYBOOK.md).
- Custom derivatives, export, Pallas, FFI: [extensions](https://skilld.dev/api/skills-raw/tristanmanchester/agent-skills/jax-development/references/ADVANCED-EXTENSIONS.md).
- Source-level diagnosis: [repo map](https://skilld.dev/api/skills-raw/tristanmanchester/agent-skills/jax-development/references/REPO-MAP.md),
  [workflow](https://skilld.dev/api/skills-raw/tristanmanchester/agent-skills/jax-development/references/EXPERT-WORKFLOW.md), and [source baseline](https://skilld.dev/api/skills-raw/tristanmanchester/agent-skills/jax-development/references/SOURCES.md).

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:

```bash
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`.