Index of every page on this site: llms.txt
All skills
tristanmanchester avatar

/jax-development

@3323bc9
by Tristan Manchestertristanmanchester/agent-skills3 stars
2
Run it yourself

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.

Use this Skill: https://skilld.dev/gh/tristanmanchester/agent-skills/jax-development

Nothing lands on disk. Nothing to clean up.

Fork this Skill

Edit a local copy. It keeps the author and licence.

referencesREPO-MAP.md

≈681 tokens on demand. Your agent reads this file only when SKILL.md points to it.

Repository map

This file assumes a local checkout similar to the provided jax-main.zip snapshot.

High-value documentation paths

Start with docs before implementation.

  • docs/installation.md
  • docs/benchmarking.md
  • docs/async_dispatch.rst
  • docs/control-flow.md
  • docs/random-numbers.md
  • docs/debugging.md
  • docs/debugging/checkify_guide.md
  • docs/debugging/print_breakpoint.md
  • docs/device_memory_profiling.md
  • docs/profiling.md
  • docs/persistent_compilation_cache.md
  • docs/buffer_donation.md
  • docs/sharded-computation.md
  • docs/notebooks/explicit-sharding.md
  • docs/migrate_pmap.md
  • docs/aot.md
  • docs/export/export.md
  • docs/export/shape_poly.md
  • docs/gpu_performance_tips.md
  • docs/faq.rst
  • docs/changelog.md
  • docs/notebooks/shard_map.md
  • docs/pallas/ and docs/jax.experimental.pallas*.rst

High-value source modules

These are the first places to inspect in the Python sources:

  • jax/_src/api.py
  • jax/_src/api_util.py
  • jax/_src/array.py
  • jax/_src/random.py
  • jax/_src/debugging.py
  • jax/_src/checkify.py
  • jax/_src/errors.py
  • jax/_src/custom_derivatives.py
  • jax/_src/ad_checkpoint.py
  • jax/_src/stages.py
  • jax/_src/export/
  • jax/_src/pjit.py
  • jax/_src/sharding.py
  • jax/_src/sharding_impls.py
  • jax/_src/mesh.py
  • jax/_src/mesh_utils.py
  • jax/_src/pallas/

Tests worth checking early

Tests are often the fastest truth source for current behaviour.

  • tests/api_test.py
  • tests/errors_test.py
  • tests/lax_control_flow_test.py
  • tests/random_test.py
  • tests/checkify_test.py
  • tests/debugging_primitives_test.py
  • tests/pjit_test.py
  • tests/pmap_test.py
  • tests/shard_map_test.py
  • tests/profiler_test.py
  • tests/export_test.py

Search recipes

Control-flow and tracer errors:

rg "ConcretizationTypeError|TracerBoolConversionError|NonConcreteBooleanIndexError" docs tests jax

Randomness:

rg "random.key|PRNGKey|key reuse|fold_in|split" docs tests jax

Sharding and migration:

rg "NamedSharding|PartitionSpec|Mesh|shard_map|pmap|migrate_pmap" docs tests jax

Debugging and profiling:

rg "debug.print|checkify|block_until_ready|profiler|compiler_ir" docs tests jax

Custom autodiff / export:

rg "custom_vjp|custom_jvp|export|ShapeDtypeStruct|shape polymorphism" docs tests jax

Pallas:

rg "pallas|mosaic|triton" docs tests jax

Recommended protocol for repo questions

  1. read the relevant docs page
  2. inspect the nearest _src implementation
  3. confirm with a nearby test
  4. only then make a behaviour claim

If docs and code look out of sync, prefer current tests plus changelog or migration docs.

Source: SKILL.md on GitHub

Skill behaviors

skilld matched fixed text patterns in SKILL.md and file names. Patterns miss obfuscated code.

skilld run checks every file with the same patterns. It asks for approval before it loads a Skill with a behavior marked Needs approval.

Third-party checks

No alerts1mo4 checks · Risk SAFE
  • Gen Agent Trust Hub1mo

    The skill provides a suite of JAX development tools. It includes functionality for dynamic loading of local modules for benchmarking and a code scanner that processes user files, which introduces a low risk of dynamic execution and indirect prompt injection.

  • Socket1mo

    No alerts

  • Snyk1mo

    Risk: LOW · No issues

  • ZeroLeaks6mo

    Score: 93/100 · 2 sections analyzed

Provenance

Signed by skilld at 3323bc9. This ties the file your Agent reads to that commit on GitHub. It does not review the instructions.

Last checked against GitHub 2 weeks ago.

Steadyupdated last month

Capability

Other metadata
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"
}

README badge

README badge for tristanmanchester/agent-skills/jax-development