mirror of
https://github.com/tristanmanchester/agent-skills.git
synced 2026-09-14 17:42:18 +08:00
Modernise jax-development: current provenance, donation-safe benchmarks, and private diagnostics (#19)
* Modernise JAX baseline and fix donated-buffer benchmarking and diagnostic disclosure * Keep JAX metadata failures structured and retain private NVIDIA diagnostics
This commit is contained in:
Tristan Manchester
committed by
GitHub
GitHub
parent
5ee544d7b1
commit
e1aa9a63f4
+88
-210
@@ -1,228 +1,106 @@
|
||||
---
|
||||
name: jax-development
|
||||
description: Use this skill when the user is writing, debugging, profiling, refactoring, reviewing, benchmarking, parallelising, exporting, or explaining JAX code, or when they mention JAX, jax.numpy, jit, grad, value_and_grad, vmap, scan, lax, random keys, pytrees, jax.Array, sharding, Mesh, PartitionSpec, NamedSharding, pmap, shard_map, Pallas, XLA, StableHLO, checkify, profiler, or the JAX repo. It helps turn NumPy or PyTorch-style code into pure functional JAX, fix tracer/control-flow/shape/PRNG bugs, remove recompiles and host-device syncs, choose transforms and sharding strategies, inspect jaxpr/lowering/IR, and benchmark compiled code correctly.
|
||||
compatibility: Best with Python 3.10+ and a working JAX installation. All bundled scripts are non-interactive and standard-library-only; they become richer when `jax` and `jaxlib` are importable or when a local JAX checkout is available.
|
||||
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:
|
||||
author: OpenAI
|
||||
version: "2.0.0"
|
||||
category: scientific-computing
|
||||
source_snapshot: "jax-main.zip from 2026-03-24"
|
||||
version: "3.0.0"
|
||||
reviewed: "2026-09-13"
|
||||
source: "https://github.com/jax-ml/jax/tree/jax-v0.11.1"
|
||||
---
|
||||
|
||||
# JAX Development
|
||||
# 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.
|
||||
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.
|
||||
|
||||
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.
|
||||
## Inspect, reproduce, measure
|
||||
|
||||
## Core promise
|
||||
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.
|
||||
|
||||
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:
|
||||
|
||||
- translating NumPy, SciPy, TensorFlow, or PyTorch code into idiomatic JAX
|
||||
- fixing tracer, control-flow, PRNG, shape, dtype, or side-effect bugs
|
||||
- choosing between `jit`, `vmap`, `scan`, `fori_loop`, `while_loop`, `cond`, `grad`, `jacrev`, `jacfwd`, `remat`, `shard_map`, or export
|
||||
- removing recompiles, host-device round trips, Python overhead, or dishonest benchmarking
|
||||
- reasoning about `jax.Array`, meshes, `PartitionSpec`, `NamedSharding`, explicit sharding, `pmap` migration, multi-host semantics, or collectives
|
||||
- using `jax.debug.print`, `checkify`, `make_jaxpr`, lowering, compiler IR, profiler traces, or memory profiling
|
||||
- using custom derivatives, export, AOT lowering, custom partitioning, Pallas, or the JAX source tree
|
||||
|
||||
Compose this skill with framework-specific skills when needed, but let this one own the JAX-specific reasoning.
|
||||
|
||||
## Do not over-apply the skill
|
||||
|
||||
Do not force JAX when the real problem is one of these instead:
|
||||
|
||||
- pure NumPy optimisation where JAX is explicitly out of scope
|
||||
- generic CUDA, Triton, NCCL, or driver debugging with no meaningful JAX component
|
||||
- framework-only design questions whose hard part is not JAX
|
||||
- irregular dynamic object-heavy Python where the right answer is probably to keep the hot path outside JAX
|
||||
|
||||
When in doubt, ask: “Is the root of the problem tracing, transformations, array semantics, compilation, sharding, or the JAX runtime?” If yes, use this skill.
|
||||
|
||||
## First-response workflow
|
||||
|
||||
### 1. Classify the task
|
||||
|
||||
Put the request into one or more lanes immediately:
|
||||
|
||||
- code design or porting
|
||||
- debugging or correctness
|
||||
- performance or compilation
|
||||
- sharding or distributed execution
|
||||
- advanced extension points
|
||||
- JAX repo navigation or source-level questions
|
||||
|
||||
Then open the matching reference file:
|
||||
|
||||
- `references/EXPERT-WORKFLOW.md` for the overall workflow
|
||||
- `references/MENTAL-MODEL.md` for tracing and staging semantics
|
||||
- `references/TRANSFORM-DECISION-MATRIX.md` for choosing primitives
|
||||
- `references/PORTING-PATTERNS.md` for NumPy or PyTorch rewrites
|
||||
- `references/CODE-REVIEW-RUBRIC.md` for self-review before replying
|
||||
- `references/DEBUGGING-TRIAGE.md` for error diagnosis
|
||||
- `references/PERFORMANCE-PLAYBOOK.md` for speed, memory, and compile-time work
|
||||
- `references/SHARDING-PLAYBOOK.md` for distributed and multi-device design
|
||||
- `references/ADVANCED-EXTENSIONS.md` for custom autodiff, export, Pallas, FFI, and internals
|
||||
- `references/REPO-MAP.md` for local source-tree navigation
|
||||
- `references/SOURCES.md` for provenance and maintenance notes
|
||||
|
||||
### 2. Inspect before guessing
|
||||
|
||||
If the problem could be environment-, backend-, or project-specific, inspect first.
|
||||
|
||||
Environment:
|
||||
```bash
|
||||
python3 scripts/jax_env_report.py --format json
|
||||
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
|
||||
```
|
||||
|
||||
Static project scan:
|
||||
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](references/MENTAL-MODEL.md),
|
||||
[transform decisions](references/TRANSFORM-DECISION-MATRIX.md), and
|
||||
[porting patterns](references/PORTING-PATTERNS.md).
|
||||
- Correctness: [debugging](references/DEBUGGING-TRIAGE.md) and
|
||||
[review rubric](references/CODE-REVIEW-RUBRIC.md).
|
||||
- Timing, compilation, memory: [performance](references/PERFORMANCE-PLAYBOOK.md).
|
||||
- Placement, collectives, multi-host: [sharding](references/SHARDING-PLAYBOOK.md).
|
||||
- Custom derivatives, export, Pallas, FFI: [extensions](references/ADVANCED-EXTENSIONS.md).
|
||||
- Source-level diagnosis: [repo map](references/REPO-MAP.md),
|
||||
[workflow](references/EXPERT-WORKFLOW.md), and [source baseline](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
|
||||
python3 scripts/jax_project_scan.py PATH --format json
|
||||
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
|
||||
```
|
||||
|
||||
Benchmark a callable honestly:
|
||||
```bash
|
||||
python3 scripts/jax_benchmark_harness.py --help
|
||||
```
|
||||
`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.
|
||||
|
||||
Inspect jaxpr, lowering, and IR:
|
||||
```bash
|
||||
python3 scripts/jax_compile_probe.py --help
|
||||
```
|
||||
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.
|
||||
|
||||
Check likely recompile behaviour across cases:
|
||||
```bash
|
||||
python3 scripts/jax_recompile_explorer.py --help
|
||||
```
|
||||
## Finish with evidence
|
||||
|
||||
Search a local JAX checkout:
|
||||
```bash
|
||||
python3 scripts/jax_repo_locator.py --help
|
||||
```
|
||||
|
||||
### 3. Reduce to a minimal reproducer
|
||||
|
||||
Prefer the smallest function that still exhibits the behaviour. JAX problems get much easier once shapes, dtypes, batching axes, randomness, and transformation boundaries are explicit.
|
||||
|
||||
### 4. Choose the least powerful mechanism that solves the problem
|
||||
|
||||
Default ordering:
|
||||
|
||||
- pure eager `jax.numpy` first
|
||||
- then `jit` or `value_and_grad`
|
||||
- then `vmap` or `scan`
|
||||
- then explicit sharding
|
||||
- then `shard_map`
|
||||
- then custom derivative, export, custom partitioning, or Pallas
|
||||
- then FFI or JAX internals
|
||||
|
||||
Escalate only with evidence.
|
||||
|
||||
### 5. End with a high-signal answer
|
||||
|
||||
Unless the user asked for something else, the reply should end with:
|
||||
|
||||
- diagnosis or design choice
|
||||
- corrected code or patch
|
||||
- why it works in JAX terms
|
||||
- how to verify it
|
||||
- remaining risks, backend caveats, or performance unknowns
|
||||
|
||||
## Expert operating rules
|
||||
|
||||
1. **Treat JAX functions as pure.** Inputs in, outputs out. Hidden mutation, global state, or implicit randomness are usually design bugs once transforms enter the picture.
|
||||
2. **Make randomness explicit.** Thread keys through the program, split once per consumer, and return updated keys when state continues.
|
||||
3. **Keep the hot path in JAX space.** Host conversion inside transformed code is almost always a bug or a sync point.
|
||||
4. **Separate static and dynamic values.** Shapes, dtypes, Python objects, and some configuration values influence tracing and compilation.
|
||||
5. **Use structured control flow.** If a branch or loop depends on array values, use JAX control-flow primitives instead of Python.
|
||||
6. **Benchmark honestly.** Warm up, block, and distinguish transfer cost, compile cost, and steady-state execution.
|
||||
7. **Optimise after evidence.** Use scans, compile probes, profiler traces, or lowering inspection before proposing deep rewrites.
|
||||
8. **Prefer current JAX idioms.** Typed keys, `jax.Array`, and modern sharding APIs are the default unless the codebase is intentionally legacy.
|
||||
9. **Think globally for sharding first.** Start with global-view code and explicit placement before dropping to per-device manual code.
|
||||
10. **Never bluff backend-specific behaviour.** CPU, GPU, TPU, and multi-host runs differ materially. Say what was verified and what was inferred.
|
||||
|
||||
## Default red flags to proactively check
|
||||
|
||||
Always scan for these, even if the user did not mention them:
|
||||
|
||||
- `np.asarray`, `.item()`, `.tolist()`, `jax.device_get`, or printing arrays in a hot path
|
||||
- Python `if`, `for`, or `while` inside transformed code
|
||||
- shape construction or indexing based on traced values
|
||||
- global or reused PRNG keys
|
||||
- repeated creation of jitted callables inside loops
|
||||
- changing shapes, dtypes, or static arguments causing compile storms
|
||||
- very large Python loops that should be `scan` or `fori_loop`
|
||||
- `pmap` code that may be better expressed with modern sharding APIs
|
||||
- unexplained precision assumptions or implicit `x64` expectations
|
||||
- replicated-versus-sharded confusion in distributed code
|
||||
|
||||
## Available scripts
|
||||
|
||||
- `scripts/jax_env_report.py` — report versions, backend, devices, config, env vars, and an optional smoke test.
|
||||
- `scripts/jax_project_scan.py` — AST-based scan for common JAX sharp bits and migration targets.
|
||||
- `scripts/jax_benchmark_harness.py` — benchmark a callable with warm-up, blocking, optional `jit`, and optional donation.
|
||||
- `scripts/jax_compile_probe.py` — inspect `eval_shape`, jaxpr, lowering, and compiler IR; optionally write artefacts to disk.
|
||||
- `scripts/jax_recompile_explorer.py` — run several input cases through a jitted function and flag likely recompiles or signature drift.
|
||||
- `scripts/jax_repo_locator.py` — search a local JAX checkout for relevant docs, tests, or source files by topic.
|
||||
|
||||
All scripts are non-interactive, support `--help`, and default to structured JSON output.
|
||||
|
||||
## Available assets
|
||||
|
||||
- `assets/mre_template.py` — minimal reproducible example template
|
||||
- `assets/training_step_template.py` — idiomatic compiled training step with explicit key plumbing
|
||||
- `assets/scan_template.py` — carry-state loop using `lax.scan`
|
||||
- `assets/sharding_template.py` — mesh plus `NamedSharding` starter
|
||||
- `assets/shard_map_template.py` — manual SPMD starter using `jax.shard_map`
|
||||
- `assets/benchmark_template.py` — honest timing pattern with warm-up and blocking
|
||||
- `assets/profile_template.py` — trace and memory-profile starter
|
||||
- `assets/checkify_template.py` — runtime checks that survive `jit`
|
||||
- `assets/custom_vjp_template.py` — custom reverse-mode rule starter
|
||||
- `assets/export_template.py` — export and serialisation starter
|
||||
- `assets/pallas_kernel_skeleton.py` — kernel-level starting point
|
||||
- `assets/issue_report_template.md` — compact bug report / investigation template
|
||||
|
||||
## Output quality bar
|
||||
|
||||
Before sending a final answer, mentally run the code or design through `references/CODE-REVIEW-RUBRIC.md`. The answer should usually satisfy all of the following:
|
||||
|
||||
- runnable or patch-ready code
|
||||
- correct transformation and sharding semantics
|
||||
- explicit discussion of compile and runtime consequences
|
||||
- no accidental host round trips in the claimed hot path
|
||||
- no hidden PRNG or state bugs
|
||||
- an honest verification method
|
||||
|
||||
## If the task is exploratory research code
|
||||
|
||||
Prefer a staged plan:
|
||||
|
||||
1. get a correct eager version in `jax.numpy`
|
||||
2. add tests or invariants
|
||||
3. add transformations one at a time
|
||||
4. benchmark and profile
|
||||
5. only then attempt aggressive sharding or kernel work
|
||||
|
||||
This workflow beats premature `jit`/`pmap`/Pallas every time.
|
||||
|
||||
## Skill maintenance
|
||||
|
||||
When updating this skill, refresh the JAX facts most likely to drift:
|
||||
|
||||
- installation guidance
|
||||
- sharding APIs and `pmap` migration status
|
||||
- randomness recommendations
|
||||
- profiler and memory-tooling guidance
|
||||
- export / AOT APIs
|
||||
- Pallas and custom extension interfaces
|
||||
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`.
|
||||
|
||||
@@ -1,168 +1,66 @@
|
||||
# Performance playbook
|
||||
|
||||
Use this file for compile-time blowups, slow steady-state execution, hidden synchronisation, memory pressure, or distributed performance work.
|
||||
## Separate the costs
|
||||
|
||||
## First rule: separate the costs
|
||||
State the workload and reference implementation, then distinguish host preparation,
|
||||
transfer, trace/lowering/compilation, dispatch, execution, materialisation, and
|
||||
communication. Validate numerical outputs and gradients before optimising. Record
|
||||
backend/devices, versions, shapes, dtypes, precision, sharding, and cache state.
|
||||
|
||||
Never talk about “JAX performance” as one number. Separate:
|
||||
## Benchmark independent inputs
|
||||
|
||||
- host-to-device transfer
|
||||
- first-call trace and compile
|
||||
- steady-state device execution
|
||||
- synchronisation / materialisation
|
||||
- communication or resharding
|
||||
- memory pressure / OOM behaviour
|
||||
The bundled `scripts/jax_benchmark_harness.py` imports a trusted module exposing
|
||||
a function and a zero-argument input factory. Example target module:
|
||||
|
||||
Most bad optimisation advice comes from mixing these.
|
||||
```python
|
||||
import jax.numpy as jnp
|
||||
|
||||
## Honest benchmark pattern
|
||||
def step(x):
|
||||
return x + 1
|
||||
|
||||
Use `assets/benchmark_template.py` or `scripts/jax_benchmark_harness.py`.
|
||||
def make_inputs():
|
||||
return (jnp.arange(1024, dtype=jnp.float32),), {}
|
||||
```
|
||||
|
||||
Checklist:
|
||||
- warm up first
|
||||
- block before stopping the timer
|
||||
- report first-call and steady-state separately
|
||||
- say whether data transfer is inside or outside the timer
|
||||
- state backend and key shapes/dtypes
|
||||
Use `--file PATH --function step --factory make_inputs --jit --repeat 20`.
|
||||
The factory must return fresh equivalent buffers for every invocation, including
|
||||
warm-up and first call. This makes explicit `--donate-argnums 0` possible without
|
||||
reusing invalidated arrays. Static arguments must remain stable; setup itself
|
||||
must not dominate memory or change the workload. The helper does not detect every
|
||||
aliasing mistake in user factories.
|
||||
|
||||
## Compile-time problems
|
||||
Inputs are synchronised outside timing; results are synchronised inside. Any
|
||||
failure propagates rather than returning a plausible timing. Report median,
|
||||
spread and raw samples, not only the fastest run. First-call time mixes multiple
|
||||
costs and can hit caches. Use lower/compile probes or separate controlled processes
|
||||
when isolated compilation/cold-cache measurements matter. A pre-jitted callable
|
||||
can consume donated inputs even without the harness's `--jit` option.
|
||||
|
||||
Symptoms:
|
||||
- first call is extremely slow
|
||||
- every call looks like a first call
|
||||
- `make_jaxpr` is huge
|
||||
- CPU-side time dominates
|
||||
A training loop is a different workload: carry the returned state into the next
|
||||
iteration, use fresh batches/keys, and separate warm-up from steady-state throughput.
|
||||
Do not keep resetting it and label the independent-input result training throughput.
|
||||
For eager/JIT comparison, use separate fresh-input runs and compare correctness
|
||||
and workload equivalence; donation invalidates shared comparison inputs too.
|
||||
|
||||
Check:
|
||||
- changing shape or dtype
|
||||
- changing sharding
|
||||
- static arguments changing every call
|
||||
- creating new lambdas/partials/jitted functions inside loops
|
||||
- long Python loops inside `jit`
|
||||
- giant captured constants / closures
|
||||
## Diagnose before rewriting
|
||||
|
||||
Tools:
|
||||
- `scripts/jax_compile_probe.py`
|
||||
- `scripts/jax_recompile_explorer.py`
|
||||
Repeated compilations: inspect changing shape/dtype/sharding/static values,
|
||||
new closures/jitted functions, large captured constants and Python-loop unrolling.
|
||||
Use the compile probe/recompile explorer as diagnostic aids and confirm with
|
||||
compilation logs/traces; a signature heuristic is not an authoritative compile count.
|
||||
|
||||
Typical fixes:
|
||||
- stabilise shapes
|
||||
- hoist `jit`
|
||||
- use `scan` or `fori_loop`
|
||||
- make static args explicit and small
|
||||
- avoid rebuilding objects each call
|
||||
Slow execution: inspect host conversions, tiny kernel dispatches, callbacks,
|
||||
input starvation, resharding, and redundant computation. Consider larger compiled
|
||||
regions, batching and structured loops only against an actual profile.
|
||||
|
||||
## Steady-state execution problems
|
||||
Memory pressure: measure live buffers and peaks, examine retained outputs,
|
||||
replication, intermediates and autodiff residuals. Donation and checkpointing
|
||||
have semantics and compute/memory trade-offs; neither is an automatic speed-up.
|
||||
Use `jax.checkpoint` before experimental alternatives unless evidence favours them.
|
||||
|
||||
Symptoms:
|
||||
- first call is fine, repeated calls are still slow
|
||||
- GPU/TPU utilisation is poor
|
||||
- code is fast in theory but not in practice
|
||||
Synchronise trace windows, capture realistic iterations, and use the target
|
||||
accelerator's profiler. CPU smoke tests cannot establish GPU/TPU throughput or
|
||||
multi-host correctness. A persistent cache helps repeated compilation but must
|
||||
be trusted and its warm/cold status included in measurements.
|
||||
|
||||
Check:
|
||||
- host round-trips
|
||||
- tiny compiled kernels separated by Python
|
||||
- poor batching
|
||||
- accidental replication or resharding
|
||||
- heavy callbacks or printing
|
||||
- slow data pipeline starving the device
|
||||
|
||||
Typical fixes:
|
||||
- larger compiled regions
|
||||
- `vmap` / `scan`
|
||||
- keep arrays on device
|
||||
- explicit sharding
|
||||
- reduce logging / callbacks in the hot path
|
||||
|
||||
## Memory problems
|
||||
|
||||
Symptoms:
|
||||
- OOM
|
||||
- high peak memory
|
||||
- code only works with tiny batch sizes
|
||||
|
||||
Check:
|
||||
- large intermediates being materialised
|
||||
- duplication across branches or batches
|
||||
- unnecessary outputs retained
|
||||
- no donation where it would help
|
||||
- sharding causing replication
|
||||
- activations that could be recomputed
|
||||
|
||||
Typical fixes:
|
||||
- buffer donation
|
||||
- rematerialisation (`jax.checkpoint`)
|
||||
- better sharding
|
||||
- smaller live ranges
|
||||
- structured loops instead of unrolled Python
|
||||
|
||||
Do not suggest donation or remat automatically; tie them to evidence.
|
||||
|
||||
## Donation
|
||||
|
||||
Donation is useful when:
|
||||
- an input buffer is dead after the call
|
||||
- the output can reuse its storage
|
||||
- memory pressure is real
|
||||
|
||||
Donation is not a magic speed-up knob. Use it primarily for memory and only after correctness and API semantics are clear.
|
||||
|
||||
## Persistent compilation cache
|
||||
|
||||
Consider it when:
|
||||
- the same program is compiled repeatedly across runs
|
||||
- compile time is a real user pain point
|
||||
- the environment is stable enough for cache reuse to matter
|
||||
|
||||
This is especially relevant for development loops and repeated workloads, not as a first response to every slowdown.
|
||||
|
||||
## Profiling strategy
|
||||
|
||||
Use a profiler before deep optimisation when:
|
||||
- the user wants serious speed work
|
||||
- compile time is not obviously the whole story
|
||||
- memory or communication may dominate
|
||||
|
||||
Start with:
|
||||
- trace collection (`assets/profile_template.py`)
|
||||
- device memory profiling for OOM or leaks
|
||||
- lowering inspection if compile structure seems wrong
|
||||
|
||||
## Performance review questions
|
||||
|
||||
Ask these in order:
|
||||
|
||||
1. Is the code timing dispatch or actual execution?
|
||||
2. Is the slow path compile, execute, transfer, or communication?
|
||||
3. Are shapes/dtypes/shardings stable?
|
||||
4. Is there a Python loop or callback in the hot path?
|
||||
5. Is the data already on device?
|
||||
6. Is sharding aligned with the algorithm?
|
||||
7. Is the memory footprint forcing a bad design?
|
||||
|
||||
## “Fast JAX code” defaults
|
||||
|
||||
These defaults are often right:
|
||||
|
||||
- compile coarse-grained steps, not every small helper
|
||||
- batch independent work with `vmap`
|
||||
- express long loops with `scan`
|
||||
- keep hot arrays on device
|
||||
- measure with blocking
|
||||
- reduce shape churn
|
||||
- prefer global-view sharding before manual per-device code
|
||||
- use profiler traces rather than intuition for serious tuning
|
||||
|
||||
## What to report back to the user
|
||||
|
||||
When performance is discussed, try to report:
|
||||
|
||||
- backend and device count
|
||||
- input shapes and dtypes
|
||||
- first-call time
|
||||
- repeated-call summary
|
||||
- what was inside the timer
|
||||
- the most likely remaining bottleneck
|
||||
|
||||
That level of honesty is more useful than a vague “should be faster”.
|
||||
See [source baseline](SOURCES.md) for current APIs and official references.
|
||||
|
||||
@@ -1,50 +1,53 @@
|
||||
# Sources and maintenance notes
|
||||
# Source baseline and maintenance
|
||||
|
||||
This skill was rebuilt from two inputs:
|
||||
Reviewed 2026-09-13. The public release baseline is **JAX 0.11.1**, released
|
||||
2026-08-17, whose PyPI metadata requires Python **>=3.12**. There is no blanket
|
||||
upper bound in that metadata; this does not guarantee wheels/backend support for
|
||||
every future interpreter. Check jaxlib, accelerator plugin, driver, and platform
|
||||
requirements separately. Python 3.13 free-threaded support was dropped in 0.11.0;
|
||||
do not conflate it with ordinary CPython 3.13.
|
||||
|
||||
1. the provided agent-skill authoring guides
|
||||
2. current JAX documentation plus the provided `jax-main.zip` source snapshot
|
||||
Traceable upstream source: tag `jax-v0.11.1`, commit
|
||||
`2d66622450e2c8633cda2307688ef7aa294bd6eb` in `jax-ml/jax`. This supersedes the
|
||||
unidentified `jax-main.zip` provenance claim; it does not assert every bundled
|
||||
example came from that tag or that a current-version runtime test was performed.
|
||||
|
||||
## JAX topics explicitly refreshed for this version
|
||||
## Current changes relevant to the retained references
|
||||
|
||||
- installation and platform guidance
|
||||
- asynchronous dispatch and honest benchmarking
|
||||
- typed PRNG keys and key-reuse considerations
|
||||
- control-flow primitives and `scan` / `fori_loop`
|
||||
- modern sharding APIs and `pmap` migration
|
||||
- export / serialisation and AOT lowering
|
||||
- profiling and memory-tooling guidance
|
||||
- Pallas and advanced extension points
|
||||
- Prefer `with jax.set_mesh(mesh):`; the old Mesh context is deprecated.
|
||||
- `NamedSharding` is a placement object. Explicit sharding-in-types requires the
|
||||
appropriate mesh axis types and propagation rules; merely constructing this
|
||||
object is not enough to claim explicit-mode semantics.
|
||||
- Ordinary pure-function design remains a useful default. Public `jax.ref` offers
|
||||
explicit stateful arrays with defined effects; inspect its rules for transforms,
|
||||
autodiff, and lifetimes rather than categorically banning all state.
|
||||
- `jnp.empty`/`empty_like` are uninitialised in 0.11.0+, not zero constructors.
|
||||
- Current export deserialisation checks its compatibility window. Keep provenance,
|
||||
producer/consumer versions and actual load/execute tests; do not bypass expiry
|
||||
errors to pretend an artefact is still supported. Prefer `in_shardings_jax` and
|
||||
`out_shardings_jax` over the deprecated HLO sharding fields.
|
||||
- Many `jax.core`/`jax.interpreters` internals were removed. Use public interfaces
|
||||
and documented `jax.extend` where appropriate; match source-level debugging to
|
||||
the exact installed commit. Experimental hijax, Pallas, and custom rematerialisation
|
||||
are targeted tools, not obligatory rewrites for every numerical function.
|
||||
- A persistent compilation cache is trusted executable infrastructure. Do not
|
||||
use a cache directory writable by untrusted users or accept arbitrary exported
|
||||
executables as inert data.
|
||||
|
||||
## Maintenance checklist
|
||||
## Primary references
|
||||
|
||||
When updating the skill for a newer JAX release:
|
||||
- [Package metadata](https://pypi.org/project/jax/)
|
||||
- [Release-tag source](https://github.com/jax-ml/jax/tree/jax-v0.11.1)
|
||||
- [Changelog](https://docs.jax.dev/en/latest/changelog.html)
|
||||
- [Installation](https://docs.jax.dev/en/latest/installation.html)
|
||||
- [Donation](https://docs.jax.dev/en/latest/buffer_donation.html)
|
||||
- [Benchmarking](https://docs.jax.dev/en/latest/benchmarking.html)
|
||||
- [Export](https://docs.jax.dev/en/latest/export/export.html)
|
||||
- [Refs](https://docs.jax.dev/en/latest/array_refs.html)
|
||||
- [Compilation cache](https://docs.jax.dev/en/latest/persistent_compilation_cache.html)
|
||||
|
||||
1. re-check:
|
||||
- `docs/changelog.md`
|
||||
- `docs/installation.md`
|
||||
- `docs/random-numbers.md`
|
||||
- `docs/debugging.md`
|
||||
- `docs/benchmarking.md`
|
||||
- `docs/sharded-computation.md`
|
||||
- `docs/migrate_pmap.md`
|
||||
- `docs/export/export.md`
|
||||
- `docs/device_memory_profiling.md`
|
||||
|
||||
2. revisit:
|
||||
- `jax/_src/api.py`
|
||||
- `jax/_src/random.py`
|
||||
- `jax/_src/debugging.py`
|
||||
- `jax/_src/pjit.py`
|
||||
- `jax/_src/sharding.py`
|
||||
- `jax/_src/pallas/`
|
||||
|
||||
3. refresh the eval prompts if terminology or recommended APIs shift
|
||||
|
||||
## Notes for future editors
|
||||
|
||||
- keep `SKILL.md` focused on workflow and escalation logic
|
||||
- push deep detail into the reference files
|
||||
- prefer scripts that produce structured output and avoid interactive prompts
|
||||
- keep claims about performance and backend behaviour tied to evidence
|
||||
- treat `pmap`, export, and Pallas guidance as likely to drift over time
|
||||
For maintenance, identify the latest stable release, read its versioned changes,
|
||||
then inspect the exact changed method/types. Do not copy unreleased examples into
|
||||
stable guidance. Test selected numerical templates, backend features, and helper
|
||||
failure paths, recording the runtime actually used. Keep static scan/evaluation
|
||||
fixtures separate from measured runtime evidence.
|
||||
|
||||
@@ -1,300 +1,87 @@
|
||||
\
|
||||
#!/usr/bin/env python3
|
||||
"""Benchmark a Python callable with optional JAX JIT and proper blocking."""
|
||||
|
||||
"""Benchmark fresh, independent JAX inputs. Imports and executes trusted project code."""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import importlib
|
||||
import contextlib
|
||||
import importlib.util
|
||||
import json
|
||||
from pathlib import Path
|
||||
import statistics
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
def load_module(module_name: str | None, file_path: str | None) -> Any:
|
||||
if bool(module_name) == bool(file_path):
|
||||
raise ValueError("Exactly one of --module or --file is required.")
|
||||
if module_name:
|
||||
return importlib.import_module(module_name)
|
||||
def measure(fn, make_inputs, jax, *, repeat=10, warmup=1):
|
||||
"""Factory returns (positional args, keyword args), with fresh donated buffers."""
|
||||
if type(repeat) is not int or repeat < 1 or type(warmup) is not int or warmup < 0:
|
||||
raise ValueError('repeat must be positive; warmup must be nonnegative')
|
||||
if not callable(fn) or not callable(make_inputs):
|
||||
raise TypeError('Function and input factory must be callable')
|
||||
|
||||
path = Path(file_path or "")
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"Module file not found: {path}")
|
||||
def once():
|
||||
args, kwargs = make_inputs()
|
||||
if not isinstance(args, (tuple, list)) or not isinstance(kwargs, dict):
|
||||
raise TypeError('Input factory must return (args sequence, kwargs dictionary)')
|
||||
# Preparation/transfers are excluded. Blocking errors must propagate.
|
||||
jax.block_until_ready((args, kwargs))
|
||||
start = time.perf_counter()
|
||||
result = fn(*args, **kwargs)
|
||||
jax.block_until_ready(result)
|
||||
return (time.perf_counter() - start) * 1000
|
||||
|
||||
spec = importlib.util.spec_from_file_location(path.stem, path)
|
||||
if spec is None or spec.loader is None:
|
||||
raise ImportError(f"Could not load module spec from: {path}")
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
first = once()
|
||||
for _ in range(warmup):
|
||||
once()
|
||||
samples = [once() for _ in range(repeat)]
|
||||
return {'first_call_ms': first, 'times_ms': samples,
|
||||
'median_ms': statistics.median(samples), 'min_ms': min(samples),
|
||||
'max_ms': max(samples), 'stdev_ms': statistics.pstdev(samples),
|
||||
'input_preparation_timed': False, 'output_synchronisation_timed': True}
|
||||
|
||||
|
||||
def resolve_attr(obj: Any, dotted: str) -> Any:
|
||||
current = obj
|
||||
for part in dotted.split("."):
|
||||
current = getattr(current, part)
|
||||
return current
|
||||
|
||||
|
||||
def load_json_arg(raw: str | None, file_path: str | None, default: Any) -> Any:
|
||||
if raw is not None and file_path is not None:
|
||||
raise ValueError("Choose either the inline JSON form or the file form, not both.")
|
||||
if raw is not None:
|
||||
return json.loads(raw)
|
||||
if file_path is not None:
|
||||
return json.loads(Path(file_path).read_text(encoding="utf-8"))
|
||||
return default
|
||||
|
||||
|
||||
def numeric_tree(value: Any) -> bool:
|
||||
if isinstance(value, (int, float, bool)):
|
||||
return True
|
||||
if isinstance(value, list):
|
||||
return all(numeric_tree(v) for v in value)
|
||||
return False
|
||||
|
||||
|
||||
def maybe_import_jax():
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument('--file', type=Path, required=True)
|
||||
parser.add_argument('--function', required=True)
|
||||
parser.add_argument('--factory', required=True, help='No-argument function returning fresh (args, kwargs)')
|
||||
parser.add_argument('--jit', action='store_true')
|
||||
parser.add_argument('--static-argnums', default='')
|
||||
parser.add_argument('--donate-argnums', default='')
|
||||
parser.add_argument('--repeat', type=int, default=10)
|
||||
parser.add_argument('--warmup', type=int, default=1)
|
||||
options = parser.parse_args()
|
||||
try:
|
||||
jax = importlib.import_module("jax")
|
||||
jnp = importlib.import_module("jax.numpy")
|
||||
return jax, jnp
|
||||
except Exception:
|
||||
return None, None
|
||||
|
||||
|
||||
def tree_arrayify(value: Any, jnp_module: Any) -> Any:
|
||||
if isinstance(value, dict):
|
||||
return {k: tree_arrayify(v, jnp_module) for k, v in value.items()}
|
||||
if isinstance(value, list):
|
||||
if numeric_tree(value):
|
||||
return jnp_module.array(value)
|
||||
return [tree_arrayify(v, jnp_module) for v in value]
|
||||
return value
|
||||
|
||||
|
||||
def tree_device_put(value: Any, jax_module: Any) -> Any:
|
||||
if isinstance(value, dict):
|
||||
return {k: tree_device_put(v, jax_module) for k, v in value.items()}
|
||||
if isinstance(value, list):
|
||||
return [tree_device_put(v, jax_module) for v in value]
|
||||
try:
|
||||
return jax_module.device_put(value)
|
||||
except Exception:
|
||||
return value
|
||||
|
||||
|
||||
def parse_int_tuple(raw: str | None) -> tuple[int, ...] | None:
|
||||
if raw is None or raw == "":
|
||||
return None
|
||||
return tuple(int(part.strip()) for part in raw.split(",") if part.strip())
|
||||
|
||||
|
||||
def block_until_ready(value: Any, jax_module: Any | None) -> Any:
|
||||
if hasattr(value, "block_until_ready"):
|
||||
return value.block_until_ready()
|
||||
if jax_module is not None:
|
||||
try:
|
||||
return jax_module.block_until_ready(value)
|
||||
except Exception:
|
||||
pass
|
||||
return value
|
||||
|
||||
|
||||
def run_once(fn: Any, args: list[Any], kwargs: dict[str, Any], jax_module: Any | None) -> Any:
|
||||
out = fn(*args, **kwargs)
|
||||
block_until_ready(out, jax_module)
|
||||
return out
|
||||
|
||||
|
||||
def summary(times_ms: list[float]) -> dict[str, float]:
|
||||
return {
|
||||
"mean_ms": statistics.mean(times_ms),
|
||||
"median_ms": statistics.median(times_ms),
|
||||
"min_ms": min(times_ms),
|
||||
"max_ms": max(times_ms),
|
||||
"stdev_ms": statistics.pstdev(times_ms) if len(times_ms) > 1 else 0.0,
|
||||
}
|
||||
|
||||
|
||||
def benchmark(fn: Any, args: list[Any], kwargs: dict[str, Any], repeat: int, jax_module: Any | None) -> list[float]:
|
||||
times_ms = []
|
||||
for _ in range(repeat):
|
||||
t0 = time.perf_counter()
|
||||
run_once(fn, args, kwargs, jax_module)
|
||||
times_ms.append((time.perf_counter() - t0) * 1e3)
|
||||
return times_ms
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Benchmark a Python callable with optional JAX JIT and proper blocking.",
|
||||
formatter_class=argparse.RawTextHelpFormatter,
|
||||
epilog="""Exit codes:
|
||||
0 success
|
||||
2 operational error
|
||||
|
||||
Examples:
|
||||
python3 scripts/jax_benchmark_harness.py --file evals/files/naive_benchmark.py --function matmul_step \\
|
||||
--args-json '[[[1.0, 2.0], [3.0, 4.0]], [[1.0], [2.0]]]' --arrayify --jit --compare-eager
|
||||
|
||||
python3 scripts/jax_benchmark_harness.py --module mypkg.train --function step \\
|
||||
--args-file args.json --kwargs-file kwargs.json --jit --repeat 20
|
||||
|
||||
python3 scripts/jax_benchmark_harness.py --file train.py --function step \\
|
||||
--args-file args.json --jit --static-argnums 2 --donate-argnums 0,1
|
||||
""",
|
||||
)
|
||||
source = parser.add_mutually_exclusive_group(required=True)
|
||||
source.add_argument("--module", help="Import path for the module containing the callable.")
|
||||
source.add_argument("--file", help="Path to a Python file containing the callable.")
|
||||
parser.add_argument("--function", required=True, help="Callable name or dotted attribute path.")
|
||||
parser.add_argument("--args-json", help="JSON list of positional arguments.")
|
||||
parser.add_argument("--args-file", help="Path to a JSON file containing positional arguments.")
|
||||
parser.add_argument("--kwargs-json", help="JSON object of keyword arguments.")
|
||||
parser.add_argument("--kwargs-file", help="Path to a JSON file containing keyword arguments.")
|
||||
parser.add_argument("--arrayify", action="store_true", help="Convert numeric JSON lists to `jax.numpy.array` when JAX is available.")
|
||||
parser.add_argument("--device-put", action="store_true", help="Apply `jax.device_put` to arguments before timing when JAX is available.")
|
||||
parser.add_argument("--jit", action="store_true", help="Wrap the callable with `jax.jit`.")
|
||||
parser.add_argument("--static-argnums", help="Comma-separated positional indices to mark static when using --jit.")
|
||||
parser.add_argument("--donate-argnums", help="Comma-separated positional indices to donate when using --jit.")
|
||||
parser.add_argument("--compare-eager", action="store_true", help="Also benchmark the original eager callable.")
|
||||
parser.add_argument("--repeat", type=int, default=10, help="Number of timed steady-state repetitions. Default: 10")
|
||||
parser.add_argument("--warmup", type=int, default=1, help="Warm-up calls before timed loops. Default: 1")
|
||||
parser.add_argument("--format", choices=("json", "text"), default="json", help="Output format. Default: json")
|
||||
parser.add_argument("--output", help="Write the report to this file instead of stdout.")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> int:
|
||||
args = parse_args()
|
||||
|
||||
try:
|
||||
module = load_module(args.module, args.file)
|
||||
fn = resolve_attr(module, args.function)
|
||||
if not callable(fn):
|
||||
raise TypeError(f"Resolved object is not callable: {args.function}")
|
||||
|
||||
raw_args = load_json_arg(args.args_json, args.args_file, [])
|
||||
raw_kwargs = load_json_arg(args.kwargs_json, args.kwargs_file, {})
|
||||
if not isinstance(raw_args, list):
|
||||
raise TypeError("Positional arguments JSON must decode to a list.")
|
||||
if not isinstance(raw_kwargs, dict):
|
||||
raise TypeError("Keyword arguments JSON must decode to an object.")
|
||||
|
||||
jax, jnp = maybe_import_jax()
|
||||
|
||||
proc_args = list(raw_args)
|
||||
proc_kwargs = dict(raw_kwargs)
|
||||
if args.arrayify:
|
||||
if jax is None or jnp is None:
|
||||
raise RuntimeError("--arrayify requires JAX to be importable.")
|
||||
proc_args = [tree_arrayify(v, jnp) for v in proc_args]
|
||||
proc_kwargs = {k: tree_arrayify(v, jnp) for k, v in proc_kwargs.items()}
|
||||
|
||||
if args.device_put:
|
||||
if jax is None:
|
||||
raise RuntimeError("--device-put requires JAX to be importable.")
|
||||
proc_args = [tree_device_put(v, jax) for v in proc_args]
|
||||
proc_kwargs = {k: tree_device_put(v, jax) for k, v in proc_kwargs.items()}
|
||||
|
||||
static_argnums = parse_int_tuple(args.static_argnums)
|
||||
donate_argnums = parse_int_tuple(args.donate_argnums)
|
||||
|
||||
report: dict[str, Any] = {
|
||||
"source": args.module or args.file,
|
||||
"callable": args.function,
|
||||
"repeat": args.repeat,
|
||||
"warmup": args.warmup,
|
||||
"jax_available": jax is not None,
|
||||
"jit_requested": args.jit,
|
||||
"compare_eager": args.compare_eager,
|
||||
"static_argnums": static_argnums,
|
||||
"donate_argnums": donate_argnums,
|
||||
}
|
||||
|
||||
if args.compare_eager or not args.jit:
|
||||
for _ in range(args.warmup):
|
||||
run_once(fn, proc_args, proc_kwargs, jax)
|
||||
eager_times = benchmark(fn, proc_args, proc_kwargs, args.repeat, jax)
|
||||
report["eager"] = {
|
||||
"times_ms": eager_times,
|
||||
"summary": summary(eager_times),
|
||||
}
|
||||
|
||||
if args.jit:
|
||||
if jax is None:
|
||||
raise RuntimeError("--jit requires JAX to be importable.")
|
||||
jit_kwargs = {}
|
||||
if static_argnums is not None:
|
||||
jit_kwargs["static_argnums"] = static_argnums
|
||||
if donate_argnums is not None:
|
||||
jit_kwargs["donate_argnums"] = donate_argnums
|
||||
fn_jit = jax.jit(fn, **jit_kwargs)
|
||||
|
||||
t0 = time.perf_counter()
|
||||
run_once(fn_jit, proc_args, proc_kwargs, jax)
|
||||
first_call_ms = (time.perf_counter() - t0) * 1e3
|
||||
|
||||
for _ in range(max(args.warmup - 1, 0)):
|
||||
run_once(fn_jit, proc_args, proc_kwargs, jax)
|
||||
|
||||
jit_times = benchmark(fn_jit, proc_args, proc_kwargs, args.repeat, jax)
|
||||
report["jit"] = {
|
||||
"first_call_ms": first_call_ms,
|
||||
"times_ms": jit_times,
|
||||
"summary": summary(jit_times),
|
||||
}
|
||||
|
||||
if "eager" in report and "jit" in report:
|
||||
eager_mean = report["eager"]["summary"]["mean_ms"]
|
||||
jit_mean = report["jit"]["summary"]["mean_ms"]
|
||||
report["speedup_vs_eager_mean"] = (eager_mean / jit_mean) if jit_mean else None
|
||||
|
||||
except Exception as exc:
|
||||
sys.stderr.write(f"Error: {type(exc).__name__}: {exc}\n")
|
||||
if not options.jit and (options.static_argnums or options.donate_argnums):
|
||||
raise ValueError('Static/donation options require --jit')
|
||||
if options.repeat < 1 or options.warmup < 0:
|
||||
raise ValueError('Invalid repeat/warmup count')
|
||||
# Keep stdout machine-readable even when the imported module prints.
|
||||
with contextlib.redirect_stdout(sys.stderr):
|
||||
import jax
|
||||
import jaxlib
|
||||
spec = importlib.util.spec_from_file_location('_jax_benchmark_target', options.file.resolve())
|
||||
if spec is None or spec.loader is None:
|
||||
raise ImportError('Cannot load benchmark module')
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[spec.name] = module
|
||||
spec.loader.exec_module(module)
|
||||
fn, factory = getattr(module, options.function), getattr(module, options.factory)
|
||||
indices = lambda text: tuple(int(x.strip()) for x in text.split(',') if x.strip())
|
||||
if options.jit:
|
||||
fn = jax.jit(fn, static_argnums=indices(options.static_argnums),
|
||||
donate_argnums=indices(options.donate_argnums))
|
||||
report = measure(fn, factory, jax, repeat=options.repeat, warmup=options.warmup)
|
||||
report.update(jax=jax.__version__, jaxlib=jaxlib.__version__,
|
||||
python=sys.version.split()[0], backend=jax.default_backend(),
|
||||
devices=[str(d) for d in jax.devices()], jit=options.jit,
|
||||
factory=options.factory, repeat=options.repeat, warmup=options.warmup,
|
||||
note='Independent-input latency, not chained training throughput or isolated compilation time')
|
||||
print(json.dumps(report, indent=2))
|
||||
return 0
|
||||
except Exception as error:
|
||||
print(json.dumps({'ok': False, 'error': f'{type(error).__name__}: {error}'}), file=sys.stderr)
|
||||
return 2
|
||||
|
||||
if args.format == "json":
|
||||
text = json.dumps(report, indent=2, sort_keys=True)
|
||||
else:
|
||||
lines = [
|
||||
f"Callable: {report['source']}::{report['callable']}",
|
||||
f"Repeat: {report['repeat']}",
|
||||
f"Warmup: {report['warmup']}",
|
||||
f"JAX available: {report['jax_available']}",
|
||||
f"JIT requested: {report['jit_requested']}",
|
||||
f"Static argnums: {report['static_argnums']}",
|
||||
f"Donate argnums: {report['donate_argnums']}",
|
||||
"",
|
||||
]
|
||||
if "eager" in report:
|
||||
lines.append("Eager")
|
||||
for key, value in report["eager"]["summary"].items():
|
||||
lines.append(f" {key}: {value:.3f}")
|
||||
lines.append("")
|
||||
if "jit" in report:
|
||||
lines.append("JIT")
|
||||
lines.append(f" first_call_ms: {report['jit']['first_call_ms']:.3f}")
|
||||
for key, value in report["jit"]["summary"].items():
|
||||
lines.append(f" {key}: {value:.3f}")
|
||||
lines.append("")
|
||||
if "speedup_vs_eager_mean" in report:
|
||||
lines.append(f"speedup_vs_eager_mean: {report['speedup_vs_eager_mean']:.3f}")
|
||||
text = "\n".join(lines)
|
||||
|
||||
if args.output:
|
||||
Path(args.output).write_text(text + ("" if text.endswith("\n") else "\n"), encoding="utf-8")
|
||||
else:
|
||||
sys.stdout.write(text)
|
||||
if not text.endswith("\n"):
|
||||
sys.stdout.write("\n")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if __name__ == '__main__':
|
||||
raise SystemExit(main())
|
||||
|
||||
@@ -1,309 +1,63 @@
|
||||
\
|
||||
#!/usr/bin/env python3
|
||||
"""Emit a structured report about the local JAX environment.
|
||||
|
||||
Design goals:
|
||||
- non-interactive
|
||||
- JSON by default
|
||||
- standard-library only
|
||||
- useful even when JAX is not importable
|
||||
"""
|
||||
|
||||
"""Report versions/devices without dumping environment values. Backend probing initialises JAX."""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import importlib
|
||||
import importlib.metadata
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import platform
|
||||
import sys
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
PACKAGE_NAMES = (
|
||||
"jax",
|
||||
"jaxlib",
|
||||
"numpy",
|
||||
"scipy",
|
||||
"flax",
|
||||
"optax",
|
||||
"equinox",
|
||||
"orbax-checkpoint",
|
||||
)
|
||||
|
||||
CONFIG_KEYS = (
|
||||
"jax_enable_x64",
|
||||
"jax_default_matmul_precision",
|
||||
"jax_debug_nans",
|
||||
"jax_debug_infs",
|
||||
"jax_debug_key_reuse",
|
||||
"jax_platform_name",
|
||||
"jax_default_prng_impl",
|
||||
"jax_transfer_guard",
|
||||
"jax_compilation_cache_dir",
|
||||
)
|
||||
|
||||
ENV_PREFIXES = (
|
||||
"JAX_",
|
||||
"XLA_",
|
||||
"CUDA_",
|
||||
"NVIDIA_",
|
||||
"ROCM",
|
||||
"HIP_",
|
||||
"TPU",
|
||||
"NCCL_",
|
||||
)
|
||||
|
||||
ENV_NAMES = {
|
||||
"CUDA_VISIBLE_DEVICES",
|
||||
"NVIDIA_VISIBLE_DEVICES",
|
||||
"XLA_FLAGS",
|
||||
"PYTHONPATH",
|
||||
"LD_LIBRARY_PATH",
|
||||
"PATH",
|
||||
}
|
||||
|
||||
|
||||
def package_info(name: str) -> dict[str, Any]:
|
||||
out: dict[str, Any] = {"installed": False}
|
||||
def build_report(smoke=False):
|
||||
packages, package_errors = {}, {}
|
||||
for name in ('jax', 'jaxlib', 'numpy', 'scipy', 'flax', 'optax', 'equinox', 'orbax-checkpoint'):
|
||||
try: packages[name] = importlib.metadata.version(name)
|
||||
except importlib.metadata.PackageNotFoundError: packages[name] = None
|
||||
except Exception as error:
|
||||
packages[name] = None
|
||||
package_errors[name] = type(error).__name__
|
||||
report = {'ok': False, 'python': platform.python_version(), 'platform': platform.platform(),
|
||||
'packages': packages, 'package_errors': package_errors, 'environment_names_only': sorted(
|
||||
k for k in os.environ if k.startswith(('JAX_', 'XLA_', 'CUDA_', 'NVIDIA_', 'NCCL_', 'TPU', 'ROCM', 'HIP_')))}
|
||||
try:
|
||||
out["version"] = importlib.metadata.version(name)
|
||||
out["installed"] = True
|
||||
except importlib.metadata.PackageNotFoundError:
|
||||
out["version"] = None
|
||||
except Exception as exc: # pragma: no cover - best effort only
|
||||
out["version"] = None
|
||||
out["error"] = f"{type(exc).__name__}: {exc}"
|
||||
return out
|
||||
|
||||
|
||||
def selected_environment() -> dict[str, str]:
|
||||
env = {}
|
||||
for key, value in os.environ.items():
|
||||
if key in ENV_NAMES or any(key.startswith(prefix) for prefix in ENV_PREFIXES):
|
||||
env[key] = value
|
||||
return dict(sorted(env.items()))
|
||||
|
||||
|
||||
def read_config_value(jax_module: Any, key: str) -> Any:
|
||||
cfg = getattr(jax_module, "config", None)
|
||||
if cfg is None:
|
||||
return None
|
||||
|
||||
# Try the most stable public-ish access patterns first.
|
||||
for getter in (
|
||||
lambda: getattr(cfg, "values", {}).get(key),
|
||||
lambda: cfg.read(key), # type: ignore[attr-defined]
|
||||
lambda: getattr(cfg, key),
|
||||
):
|
||||
try:
|
||||
value = getter()
|
||||
if value is not None:
|
||||
return value
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
def maybe_key(jax_module: Any, seed: int):
|
||||
try:
|
||||
return jax_module.random.key(seed)
|
||||
except Exception:
|
||||
return jax_module.random.PRNGKey(seed)
|
||||
|
||||
|
||||
def smoke_test(jax_module: Any, jnp_module: Any) -> dict[str, Any]:
|
||||
report: dict[str, Any] = {"ok": False}
|
||||
t0 = time.perf_counter()
|
||||
|
||||
key = maybe_key(jax_module, 0)
|
||||
x = jax_module.random.normal(key, (256, 256), dtype=jnp_module.float32)
|
||||
y = (x @ x.T).block_until_ready()
|
||||
|
||||
@jax_module.jit
|
||||
def loss_fn(z):
|
||||
return jnp_module.sum(jnp_module.tanh(z @ z.T))
|
||||
|
||||
loss = loss_fn(x)
|
||||
loss.block_until_ready()
|
||||
grad = jax_module.grad(lambda z: jnp_module.sum(jnp_module.sin(z)))(x)
|
||||
jax_module.block_until_ready(grad)
|
||||
|
||||
report["ok"] = True
|
||||
report["elapsed_ms"] = (time.perf_counter() - t0) * 1e3
|
||||
report["matmul_shape"] = tuple(int(v) for v in y.shape)
|
||||
report["loss_dtype"] = str(loss.dtype)
|
||||
import jax
|
||||
import jax.numpy as jnp
|
||||
report.update(backend=jax.default_backend(), devices=[str(d) for d in jax.devices()],
|
||||
process_count=jax.process_count(), process_index=jax.process_index(),
|
||||
x64=bool(jax.config.jax_enable_x64))
|
||||
if smoke:
|
||||
x = jnp.arange(8, dtype=jnp.float32)
|
||||
output = jax.jit(lambda a: a + 1)(x)
|
||||
jax.block_until_ready(output)
|
||||
report['smoke_test'] = {'ok': bool(jnp.all(output == x + 1)), 'scope': 'small compiled addition'}
|
||||
if not report['smoke_test']['ok']: return report
|
||||
report['ok'] = not package_errors
|
||||
except Exception as error:
|
||||
# Exception messages may contain paths, endpoint addresses, or configuration values.
|
||||
report['error_type'] = type(error).__name__
|
||||
report['next_step'] = 'Inspect the original import/backend error privately in the target environment'
|
||||
return report
|
||||
|
||||
|
||||
def load_jax_report(run_smoke_test: bool) -> dict[str, Any]:
|
||||
report: dict[str, Any] = {"imported": False}
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument('--smoke-test', action='store_true')
|
||||
parser.add_argument('--format', choices=('json', 'text'), default='json')
|
||||
parser.add_argument('--output', type=Path)
|
||||
args = parser.parse_args()
|
||||
report = build_report(args.smoke_test)
|
||||
text = json.dumps(report, indent=2) if args.format == 'json' else '\n'.join(f'{k}: {v}' for k, v in report.items())
|
||||
try:
|
||||
jax = importlib.import_module("jax")
|
||||
jnp = importlib.import_module("jax.numpy")
|
||||
except Exception as exc:
|
||||
report["import_error"] = f"{type(exc).__name__}: {exc}"
|
||||
return report
|
||||
|
||||
report["imported"] = True
|
||||
report["version"] = getattr(jax, "__version__", None)
|
||||
|
||||
try:
|
||||
report["default_backend"] = jax.default_backend()
|
||||
except Exception:
|
||||
report["default_backend"] = None
|
||||
|
||||
try:
|
||||
report["process_count"] = int(jax.process_count())
|
||||
report["process_index"] = int(jax.process_index())
|
||||
report["device_count"] = int(jax.device_count())
|
||||
report["local_device_count"] = int(jax.local_device_count())
|
||||
except Exception as exc:
|
||||
report["process_error"] = f"{type(exc).__name__}: {exc}"
|
||||
|
||||
devices = []
|
||||
try:
|
||||
for dev in jax.devices():
|
||||
devices.append(
|
||||
{
|
||||
"id": getattr(dev, "id", None),
|
||||
"platform": getattr(dev, "platform", None),
|
||||
"device_kind": getattr(dev, "device_kind", None),
|
||||
"process_index": getattr(dev, "process_index", None),
|
||||
"memory_limit": getattr(dev, "memory_limit", None),
|
||||
}
|
||||
)
|
||||
except Exception as exc:
|
||||
report["devices_error"] = f"{type(exc).__name__}: {exc}"
|
||||
report["devices"] = devices
|
||||
|
||||
cfg = {}
|
||||
for key in CONFIG_KEYS:
|
||||
cfg[key] = read_config_value(jax, key)
|
||||
report["config"] = cfg
|
||||
|
||||
if run_smoke_test:
|
||||
try:
|
||||
report["smoke_test"] = smoke_test(jax, jnp)
|
||||
except Exception as exc:
|
||||
report["smoke_test"] = {"ok": False, "error": f"{type(exc).__name__}: {exc}"}
|
||||
|
||||
return report
|
||||
|
||||
|
||||
def build_report(run_smoke_test: bool) -> dict[str, Any]:
|
||||
return {
|
||||
"python": {
|
||||
"version": sys.version.split()[0],
|
||||
"executable": sys.executable,
|
||||
},
|
||||
"platform": {
|
||||
"system": platform.system(),
|
||||
"release": platform.release(),
|
||||
"machine": platform.machine(),
|
||||
"platform": platform.platform(),
|
||||
},
|
||||
"packages": {name: package_info(name) for name in PACKAGE_NAMES},
|
||||
"environment": selected_environment(),
|
||||
"jax": load_jax_report(run_smoke_test),
|
||||
}
|
||||
|
||||
|
||||
def format_text(report: dict[str, Any]) -> str:
|
||||
lines = []
|
||||
lines.append("Python")
|
||||
lines.append(f" version: {report['python']['version']}")
|
||||
lines.append(f" executable: {report['python']['executable']}")
|
||||
lines.append("")
|
||||
lines.append("Platform")
|
||||
for key, value in report["platform"].items():
|
||||
lines.append(f" {key}: {value}")
|
||||
lines.append("")
|
||||
lines.append("Packages")
|
||||
for name, info in report["packages"].items():
|
||||
status = info.get("version") if info.get("installed") else "not installed"
|
||||
lines.append(f" {name}: {status}")
|
||||
lines.append("")
|
||||
lines.append("Environment")
|
||||
for key, value in report["environment"].items():
|
||||
lines.append(f" {key}={value}")
|
||||
lines.append("")
|
||||
|
||||
jax_info = report["jax"]
|
||||
lines.append("JAX")
|
||||
if not jax_info.get("imported"):
|
||||
lines.append(f" import_failed: {jax_info.get('import_error')}")
|
||||
return "\n".join(lines)
|
||||
|
||||
lines.append(f" version: {jax_info.get('version')}")
|
||||
lines.append(f" default_backend: {jax_info.get('default_backend')}")
|
||||
lines.append(f" process_count: {jax_info.get('process_count')}")
|
||||
lines.append(f" process_index: {jax_info.get('process_index')}")
|
||||
lines.append(f" device_count: {jax_info.get('device_count')}")
|
||||
lines.append(f" local_device_count: {jax_info.get('local_device_count')}")
|
||||
lines.append("")
|
||||
lines.append("Devices")
|
||||
for device in jax_info.get("devices", []):
|
||||
lines.append(
|
||||
" - id={id} platform={platform} kind={device_kind} process_index={process_index} memory_limit={memory_limit}".format(
|
||||
**device
|
||||
)
|
||||
)
|
||||
lines.append("")
|
||||
lines.append("Config")
|
||||
for key, value in jax_info.get("config", {}).items():
|
||||
lines.append(f" {key}: {value}")
|
||||
|
||||
smoke = jax_info.get("smoke_test")
|
||||
if smoke:
|
||||
lines.append("")
|
||||
lines.append("Smoke test")
|
||||
for key, value in smoke.items():
|
||||
lines.append(f" {key}: {value}")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Emit a structured report about the local JAX environment.",
|
||||
formatter_class=argparse.RawTextHelpFormatter,
|
||||
epilog="""Exit codes:
|
||||
0 success
|
||||
2 operational error
|
||||
|
||||
Examples:
|
||||
python3 scripts/jax_env_report.py
|
||||
python3 scripts/jax_env_report.py --smoke-test --format text
|
||||
python3 scripts/jax_env_report.py --output env.json
|
||||
""",
|
||||
)
|
||||
parser.add_argument("--smoke-test", action="store_true", help="Run a small JAX smoke test if JAX is importable.")
|
||||
parser.add_argument("--format", choices=("json", "text"), default="json", help="Output format. Default: json")
|
||||
parser.add_argument("--output", help="Write output to a file instead of stdout.")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> int:
|
||||
args = parse_args()
|
||||
try:
|
||||
report = build_report(args.smoke_test)
|
||||
text = json.dumps(report, indent=2, sort_keys=True) if args.format == "json" else format_text(report)
|
||||
if args.output:
|
||||
with open(args.output, "w", encoding="utf-8") as f:
|
||||
f.write(text)
|
||||
if not text.endswith("\n"):
|
||||
f.write("\n")
|
||||
else:
|
||||
sys.stdout.write(text)
|
||||
if not text.endswith("\n"):
|
||||
sys.stdout.write("\n")
|
||||
return 0
|
||||
except Exception as exc:
|
||||
sys.stderr.write(f"Error: {type(exc).__name__}: {exc}\n")
|
||||
descriptor = os.open(args.output, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
|
||||
with os.fdopen(descriptor, 'w', encoding='utf-8') as stream: stream.write(text + '\n')
|
||||
else: print(text)
|
||||
except OSError as error:
|
||||
print(json.dumps({'ok': False, 'error_type': type(error).__name__}), file=sys.stderr)
|
||||
return 2
|
||||
return 0 if report['ok'] else 2
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
if __name__ == '__main__': raise SystemExit(main())
|
||||
|
||||
@@ -0,0 +1,54 @@
|
||||
import importlib.util
|
||||
from pathlib import Path
|
||||
import unittest
|
||||
|
||||
spec = importlib.util.spec_from_file_location('bench', Path(__file__).parents[1] / 'scripts/jax_benchmark_harness.py')
|
||||
bench = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(bench)
|
||||
|
||||
class FakeJax:
|
||||
def __init__(self): self.blocks = []
|
||||
def block_until_ready(self, value): self.blocks.append(value)
|
||||
|
||||
class BenchmarkTests(unittest.TestCase):
|
||||
def test_fresh_inputs_for_every_call(self):
|
||||
created = []
|
||||
def factory():
|
||||
value = {'used': False}
|
||||
created.append(value)
|
||||
return ([value], {})
|
||||
def donated(value):
|
||||
self.assertFalse(value['used'])
|
||||
value['used'] = True
|
||||
return {'result': 1}
|
||||
fake = FakeJax()
|
||||
result = bench.measure(donated, factory, fake, repeat=3, warmup=2)
|
||||
self.assertEqual(len(created), 6)
|
||||
self.assertEqual(len(fake.blocks), 12)
|
||||
self.assertEqual(len(result['times_ms']), 3)
|
||||
def test_block_failure_is_not_swallowed(self):
|
||||
class Broken:
|
||||
def block_until_ready(self, value): raise RuntimeError('device failure')
|
||||
with self.assertRaisesRegex(RuntimeError, 'device failure'):
|
||||
bench.measure(lambda x: x, lambda: ([1], {}), Broken())
|
||||
def test_output_failure_propagates(self):
|
||||
class Broken:
|
||||
def block_until_ready(self, value):
|
||||
if value == 'output': raise RuntimeError('asynchronous failure')
|
||||
with self.assertRaisesRegex(RuntimeError, 'asynchronous failure'):
|
||||
bench.measure(lambda: 'output', lambda: ([], {}), Broken())
|
||||
def test_invalid_counts_fail_before_factory(self):
|
||||
for repeat, warmup in [(0, 1), (1, -1), (True, 1), (1, 1.5)]:
|
||||
with self.assertRaises(ValueError):
|
||||
bench.measure(None, None, FakeJax(), repeat=repeat, warmup=warmup)
|
||||
def test_malformed_factory_fails(self):
|
||||
with self.assertRaises(TypeError):
|
||||
bench.measure(lambda: None, lambda: ('bad', {}), FakeJax())
|
||||
def test_actual_cpu_donation(self):
|
||||
import jax
|
||||
import jax.numpy as jnp
|
||||
fn = jax.jit(lambda x: x + 1, donate_argnums=(0,))
|
||||
result = bench.measure(fn, lambda: ([jnp.arange(8, dtype=jnp.float32)], {}), jax, repeat=3)
|
||||
self.assertEqual(len(result['times_ms']), 3)
|
||||
|
||||
if __name__ == '__main__': unittest.main()
|
||||
@@ -0,0 +1,22 @@
|
||||
import importlib.util
|
||||
import json
|
||||
from pathlib import Path
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
spec = importlib.util.spec_from_file_location('env_report', Path(__file__).parents[1] / 'scripts/jax_env_report.py')
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
|
||||
class EnvironmentTests(unittest.TestCase):
|
||||
def test_secret_values_not_emitted(self):
|
||||
with patch.dict('os.environ', {'JAX_TEST_TOKEN':'not-for-the-report', 'XLA_TEST_PATH':'/private/test-path'}):
|
||||
report = module.build_report()
|
||||
self.assertNotIn('not-for-the-report', json.dumps(report))
|
||||
self.assertNotIn('/private/test-path', json.dumps(report))
|
||||
self.assertIn('JAX_TEST_TOKEN', report['environment_names_only'])
|
||||
def test_cpu_smoke(self):
|
||||
report = module.build_report(True)
|
||||
self.assertTrue(report['ok'])
|
||||
self.assertTrue(report['smoke_test']['ok'])
|
||||
|
||||
if __name__ == '__main__': unittest.main()
|
||||
@@ -0,0 +1,16 @@
|
||||
import json
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
from test_env import module
|
||||
|
||||
class EnvironmentReviewTests(unittest.TestCase):
|
||||
def test_nvidia_names_without_values(self):
|
||||
with patch.dict('os.environ',{'NVIDIA_VISIBLE_DEVICES':'secret-device-set'}):
|
||||
report=module.build_report()
|
||||
self.assertIn('NVIDIA_VISIBLE_DEVICES',report['environment_names_only'])
|
||||
self.assertNotIn('secret-device-set',json.dumps(report))
|
||||
def test_unreadable_package_metadata_is_structured_and_not_ok(self):
|
||||
with patch.object(module.importlib.metadata,'version',side_effect=PermissionError('private-message')):
|
||||
report=module.build_report()
|
||||
self.assertFalse(report['ok']);self.assertEqual(report['package_errors']['jax'],'PermissionError')
|
||||
self.assertNotIn('private-message',json.dumps(report))
|
||||
Reference in New Issue
Block a user