dev-jax

Guides writing, transforming, and profiling JAX code with explicit PRNG keys and pytrees.

Updated Jul 3, 2026

Search Tags:#profiling #jax #jit #vmap #prng #pytrees #optax

One-click install

npx skills add https://github.com/mnazaal/dotfiles --skill dev-jax-mnazaal

Or copy as Structured Prompt for Agent▼

Please help me install this Agent Skill.
Skill: dev-jax
Source: https://github.com/mnazaal/dotfiles/tree/main/.agents/skills/dev-jax
Command: npx skills add https://github.com/mnazaal/dotfiles --skill dev-jax-mnazaal

SYSTEM DOCUMENTATION & REQUIREMENTS

What problem does it solve? JAX code fails in non-obvious ways: stale APIs from outdated training data, silent retracing that makes training 100x slower, NaN gradients from jnp.where on singular functions, and OOM jobs with no clear cause. This Skill encodes the idioms, anti-patterns, and current API surface needed to write correct and fast JAX programs. ## Core Features & Use Cases - JAX Idioms & Transformations: Enforces pure functions, explicit PRNG key splitting, pytree discipline, and correct use of jit, vmap, scan, and grad. - Stale API Detection: Maintains a verified table of removed or renamed APIs (jax.tree_map, PRNGKey, jnp.trapz, jax.experimental.optimizers) against the installed JAX version. - Performance & Memory Diagnosis: Provides an ordered debugging workflow for slow or OOMing jobs, covering retracing, dispatch overhead, and allocation measurement before tuning. - Use Case: A training loop calls bare jax.grad plus optax.update in a Python for loop and runs 100x slower than expected; the Skill directs you to jit the combined grad-and-update step and verify with a cache-size check. ## Quick Start Use the dev-jax skill to review my JAX training script for retracing issues and outdated API calls.

Frequently Asked Questions about dev-jax

High-intent search queries and answers about installing and using this skill.

FAQPage Schema

How do I debug an out-of-memory JAX job?▼

Measure what actually allocates before shrinking anything, and treat batch size as the last knob rather than the first. For slow jobs, sync with block_until_ready before timing, discard compile steps from averages, and check GPU utilization, input bottlenecks, and retracing in that order.