python-jax
CommunityHigh-performance NumPy-like computing with JAX.
Authorjkitchin
Version1.0.0
Installs0
System Documentation
What problem does it solve?
Expert guidance for JAX, enabling accelerator-focused array computation and program transformations such as automatic differentiation, JIT compilation, vectorization, and parallelization.
Core Features & Use Cases
- Transformations: automatic differentiation with
jax.grad, JIT compilation withjax.jit, automatic vectorization withjax.vmap, and parallelization withjax.pmap. - Differentiation & Transformations: support for higher-order derivatives and nested pytrees.
- Performance & Hardware: seamless GPU/TPU acceleration with minimal code changes.
- Use cases include custom gradient-based optimization, high-performance numerical kernels, and scalable ML research workflows.
Quick Start
Create a small function, differentiate it with jax.grad, and run the gradient on a sample vector to observe automatic differentiation in action.
Dependency Matrix
Required Modules
None requiredComponents
Standard package💻 Claude Code Installation
Recommended: Let Claude install automatically. Simply copy and paste the text below to Claude Code.
Please help me install this Skill: Name: python-jax Download link: https://github.com/jkitchin/skillz/archive/main.zip#python-jax Please download this .zip file, extract it, and install it in the .claude/skills/ directory.