JAX
- Repo stars 39
- Author repo awesome-omni-skill
JAX Skill
JAX is Autograd and XLA, brought together for high-performance machine learning research.
Contents
- Concepts & Theory
- Immutability
- The 4 Transformations
- Pytrees
- Code Examples
jit,grad,vmap,randomusage- Control Flow (
scan,cond,fori_loop) - Parallelism (
sharding)
Common Workflows
1. Developing a new Model
- Define your parameters as a Pytree (dict/dataclass).
- Define your forward pass function (pure).
- Define your loss function.
- Use
jax.value_and_gradto get gradients. - Use
jax.jitto speed up the update step. - See examples.md for snippets.
2. Debugging Shapes/NaNs
- Disable JIT:
jax.config.update("jax_disable_jit", True)to debug with standard python tools. - Use
jax.debug.printinside JITted functions.
- Fluxly category
- AI
- Author-declared agents
- No explicit declaration found; this is not inferred or tested compatibility
- Static check
- 88 / 100 · heuristic scan, not runtime safety proof
- Author / version / license
- @diegosouzapw · no license declared
- Fluxly token estimate
- Lean
- Fluxly setup estimate
- Plug-and-play
- External API key
- No requirement detected
- Detected OS requirements
- Unspecified
- Runtime requirements
- Python
- Detected file/system behavior
-
- Read-only
- Write / modify
- Detected network behavior
- Local-only
- Install commands
- None (reference only)
Profile is derived at build time from SKILL.md and install vectors. Subject to drift from author intent.
Heads up: 未限定 allowed-tools,默认拥有全部工具权限。
The current SKILL.md does not define a fixed output example. Contents
Concepts & Theory Immutability The 4 Transformations
Common Workflows
Common Workflows
1. Developing a new Model
Define your parameters as a Pytree (dict/dataclass). Define your forward pass function (pure). Define your loss function.
2. Debugging Shapes/NaNs
Disable JIT: jax.config.update("jaxdisablejit", True) to debug with standard python tools. Use jax.debug.print inside JITted functions.
# 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. Evidence boundary and execution chain
Author text anchors workflow facts; Fluxly only indexes current sections, terms, files, and commands.
sections -> Contents → Common Workflows → 1. Developing a new Model → 2. Debugging Shapes/NaNs
terms -> JAX is Autograd and XLA, brought together for high-performance machine learning research.
files/cmd -> jit · grad · vmap · random · scan · cond · foriloop · sharding
body sha256 -> 6dc651c2fe13
Decide Fit First
Design Intent
How To Use It
Boundaries And Review