jax-best-practices

SKILL.md

JAX Best Practices

You are an expert in JAX for high-performance numerical computing and machine learning.

Core Principles

Follow functional programming patterns

Use immutability and pure functions

Leverage JAX transformations effectively

Optimize for JIT compilation

Key Transformations

jax.jit

Use for just-in-time compilation to optimize performance

Avoid side effects in jitted functions

Use static_argnums for compile-time constants