jax

separate

High-performance numerical computing and machine learning workflows using JAX. Supports array operations, automatic differentiation, JIT compilation, RNN-style scans, map/reduce operations, and gradient computations. Ideal for scientific computing, ML models, and dynamic array transformations.

When to use

array operations
automatic differentiation
JIT compilation
RNN-style scans
map/reduce operations
gradient computations
scientific computing
ML models

Skill page

Source

Repository	Path
benchflow-ai/skillsbench	`tasks/jax-computing-basics/environment/skills/jax-skills/SKILL.md`