Essential tools for using JAX in machine learning and mathematical analysis, covering core concepts, transformations, ML specifics, control flow, and parallelism.
Installs into .claude/skills of the current project.
Are you the author of Jax?
Add the live security badge to your README. It updates with every re-scan.
[](https://www.skillsdirectory.com/skills/diegosouzapw-jax)
---
name: JAX
description: Essential tools for using JAX in machine learning and mathematical analysis, covering core concepts, transformations, ML specifics, control flow, and parallelism.
---
# 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.
