# JAX

> Provides guidance on using JAX for machine learning and mathematical analysis, covering core concepts, transformations, ML specifics, control flow, and parallelism.

- Skill: `diegosouzapw/jax` (Agent Skill, multi-file: 2 files)
- Install (CLI): `npx skillmds@latest add diegosouzapw/jax`
- Raw SKILL.md: https://api.skillmd.com/api/skills/diegosouzapw/jax/raw
- Safety review: PASS (external: skill-scanner PASS, skillspector PASS)
- Works with: Claude Code, Claude.ai, OpenAI Codex
- Category: AI & ML, Coding & Dev Tools, Model Training & Fine-tuning
- Tags: Autograd, Jax, Jit, Machine Learning, Parallelism, Pytrees, Vmap, Xla
- Author: diegosouzapw (https://skillmd.com/u/diegosouzapw)
- Updated: 2026-08-22
- Page: https://skillmd.com/skills/diegosouzapw/jax

---

# JAX Skill

JAX is Autograd and XLA, brought together for high-performance machine learning research.

## Contents

- [Concepts & Theory](reference.md)
  - Immutability
  - The 4 Transformations
  - Pytrees
- [Code Examples](examples.md)
  - `jit`, `grad`, `vmap`, `random` usage
  - Control Flow (`scan`, `cond`, `fori_loop`)
  - Parallelism (`sharding`)

## Common Workflows

### 1. Developing a new Model
1.  Define your parameters as a Pytree (dict/dataclass).
2.  Define your forward pass function (pure).
3.  Define your loss function.
4.  Use `jax.value_and_grad` to get gradients.
5.  Use `jax.jit` to speed up the update step.
6.  See [examples.md](examples.md) for snippets.

### 2. Debugging Shapes/NaNs
1.  Disable JIT: `jax.config.update("jax_disable_jit", True)` to debug with standard python tools.
2.  Use `jax.debug.print` inside JITted functions.