jax-patterns
JAX, Flax, Optax, and Equinox patterns for ML training. Covers JIT, vmap, pmap, TPU usage, and functional model design. Use when building or debugging JAX-based training pipelines.
npx skills add majiayu000/claude-skill-registry --skill jax-patterns --agent claude-code
Same command for any agent — swap --agent for codex, cursor, copilot.
Weekly change comes from our own snapshots, not the repository page — it measures attention, not adoption.
# JAX Patterns ## Framework Selection | Framework | Style | Strengths | Weaknesses | Use When | |-----------|-------|-----------|------------|----------| | **Flax (NNX)** | Pythonic, mutable | Google-backed, large ecosystem, NNX is latest API | Linen API is legacy; NNX still maturing | Default for new projects, TPU training | | **Equinox** | PyTorch-like, pytree-native | Clean API, composable, feels natural | Smaller ecosystem, fewer examples | Prefer PyTorch style, research code | | **Haiku** | Functional transforms | DeepMind ecosystem, simple | Maintenance mode (use Flax NNX instead) | Legacy DeepMind codebases only | Default recommendation: Flax NNX for production/TPU work. Equinox for research where API ergonomics matter. ## Core JAX Concepts ### JIT Compilation ```python import jax import jax.numpy as jnp from functools import partial # Basic JIT -- function must be pure (no side effects) @jax.jit def compute(x, y): return jnp.dot(x, y) + jnp.sin(x) # Static arguments -- use for values that change compilation (shapes, dtypes, booleans) @partial(jax.jit, static_argnums=(2,)) def forward(params, x, use_dropout: bool): ... ``` ### vmap -- Automatic Batching ```python def single_
- Framework Selection
- Core JAX Concepts
- JIT Compilation
- vmap -- Automatic Batching
- pmap -- Data Parallelism
- Flax NNX Model Definition
- JIT-Compiled Training Step (Flax NNX)
- Equinox Alternative
- Optax Optimizer Chains
- Custom Loss with Auxiliary Data
- TPU-Specific Patterns
- Gotchas and Anti-Patterns
- Functional Constraints
- Pytree Handling
What does the jax-patterns skill do?
JAX, Flax, Optax, and Equinox patterns for ML training. Covers JIT, vmap, pmap, TPU usage, and functional model design. Use when building or debugging JAX-based training pipelines.
How do I install it?
Run `npx skills add majiayu000/claude-skill-registry --skill jax-patterns --agent claude-code` — it drops the skill into your project so the agent can pick it up. Swap the --agent value for codex, cursor or copilot if you use one of those.
Where does this skill come from?
From majiayu000/claude-skill-registry, a repository with 534 stars. We read it straight from the repository tree rather than a submitted listing, so what you see here is what is actually published.
Is a popular skill a good skill?
Not necessarily. Stars measure attention, not adoption — a repository can trend for a week and be abandoned. That is why we show the weekly change from our own snapshots next to the total, instead of a single flattering number.