jax-mcp

MCP (Model Context Protocol) server for JAX documentation.

Enables LLMs to access up-to-date JAX documentation and validate generated code.

Installation

pip install -e .

Or run directly:

python -m jax_mcp

Usage with Claude Code

# Add as MCP server
claude mcp add -t stdio -s user jax -- python -m jax_mcp

# Or if installed globally
claude mcp add -t stdio -s user jax -- jax-mcp

Configuration

Offline Mode

Point to a local JAX clone for offline access:

export JAX_DOCS_PATH=/path/to/jax/docs
python -m jax_mcp

Online Mode (Default)

Fetches docs from GitHub with local caching:

python -m jax_mcp
# Fetches from: raw.githubusercontent.com/google/jax/main/docs/

Available Tools

Documentation Categories

• concepts: Core JAX concepts (pytrees, transformations, tracing)
• gotchas: Common mistakes and how to avoid them
• transforms: jit, vmap, grad, scan patterns
• advanced: Distributed computing, custom pytrees
• performance: GPU tips, profiling, benchmarking
• api: Module overviews (jax.numpy, jax.lax, jax.random)
• examples: Practical code examples

JAX Checker

The jax-checker tool catches common JAX mistakes:

• In-place array mutations (array[idx] = value)
• Side effects in jitted functions (print, globals)
• PRNG key reuse without splitting
• Python control flow in traced code
• Missing block_until_ready() for benchmarks
• Float64 usage without config

Development

# Install dev dependencies
pip install -e ".[dev]"

# Run tests
pytest

License

MIT