python-jax

Files (1)

• SKILL.md

SKILL.md (first lines)

# JAX - High-Performance Numerical Computing

## Overview

**JAX** is a Python library for accelerator-oriented array computation and program transformation, designed for high-performance numerical computing and large-scale machine learning. It combines a familiar NumPy-style API with powerful function transformations for automatic differentiation, compilation, vectorization, and parallelization.

**Core value:** Write NumPy-like Python code and automatically get gradients, GPU/TPU acceleration, vectorization, and parallelization through composable function transformations—without changing your mathematical notation.

## When to Use JAX

**Use JAX when:**

- Need automatic differentiation for optimization or machine learning
- Want GPU/TPU acceleration with minimal code changes
- Require high-performance numerical computing
- Building custom gradient-based algorithms
- Need to vectorize or parallelize functions automatically
- Working on research requiring flexible differentiation
- Want functional programming approach to numerical code

**Don't use when:**

- Simple NumPy operations without performance needs (overhead not justified)
- Heavy reliance on in-place mutations (JAX arrays are immutable)
- Imperative, stateful code with side effects
- Need control flow depending on runtime data values (limited support)
- Working with libraries that require NumPy arrays (compatibility issues)

## Core Transformations

JAX provides four fundamental, composable transformations:

### 1. `jax.grad` - Automatic Differentiation

Compute gradients automatically using reverse-mode autodiff.

```python
import jax
import jax.numpy as jnp

def loss(x):
    return jnp.sum(x**2)

# Get gradient function
grad_loss = jax.grad(loss)

x = jnp.array([1.0, 2.0, 3.0])
gradient = grad_loss(x)
print(gradient)  # [2. 4. 6.]
```

**Key features:**

- Returns a function that computes gradients
- Composes for higher-order derivatives
- Works with complex nested structures (pytrees)

### 2. `jax.jit` - Just-In-Time Compilation

Compile functions to XLA for dramatic speedups.

```python
import jax

@jax.jit
def fast_function(x):
    return jnp.sum(x**2 + 3*x + 1)

# First call: compilation + execution (slower)
result = fast_function(jnp.array([1.0, 2.0, 3.0]))

# Subsequent calls: cached compiled version (much faster)
result = fast_function(jnp.array([4.0, 5.0, 6.0]))
```

**Performance:**

- 10-100x speedup typical for numerical functions
- First call slower (compilation overhead)
- Cached for subsequent calls with same shapes/dtypes

### 3. `jax.vmap` - Automatic Vectorization

Vectorize functions across batch dimensions automatically.

```python
import jax

def process_single(x):
    """Process single example"""
    return jnp.sum(x**2)

# Vectorize to process batch
process_batch = jax.vmap(process_single)

# Works on batch automatically
batch = jnp.array([[1, 2], [3, 4], [5, 6]])
results = process_batch(batch)
print(results)  # [5 25 61]
```

**Benefits:**

- Eliminates manual loop writing
- Often faster than explicit loops
- Cleaner, more declarative code

### 4. `jax.pmap` - Parallel Map

Parallelize across multiple devices (GPUs/TPUs).

```python
import jax

@jax.pmap
def parallel_fn(x):
    return x**2

# Runs on all available devices

From woernfl/ai-lab, MIT.

More from woernfl/ai-lab