JAX Best Practices: High-Performance Python Coding for Machine Learning with JIT and Vectorization

### Context
JAX is a powerful Python library for high-performance numerical computing and machine learning, built on NumPy-like APIs with automatic differentiation and JIT compilation. These best practices ensure your code is efficient, pure, and optimized for accelerators like GPUs/TPUs. Focus on functional style, immutability, and JAX transformations to avoid pitfalls and maximize performance.

### Rules
- **Adopt Functional Programming**: Stick to pure functions without side effects or mutable state; pass all dependencies explicitly to enable transformations like `jax.jit`, `jax.grad`, and `jax.vmap`.
- **Use JAX NumPy**: Replace `numpy` with `jax.numpy` (import as `jnp`) for full JAX compatibility in array operations.
- **Enable JIT Compilation**: Wrap performance-critical functions with `jax.jit` to compile to XLA for speedups, ensuring no Python control flow or side effects inside.
- **Leverage Automatic Differentiation**: Use `jax.grad`, `jax.value_and_grad` for gradients; design functions to take array inputs and output scalars for grad computation.
- **Vectorize with vmap**: Apply `jax.vmap` to batch over axes instead of loops; prefer `jax.lax.scan` for sequential data.
- **Handle Control Flow**: Inside JIT, use `jax.lax.cond`, `jax.lax.fori_loop`, or `jax.lax.while_loop` instead of Python `if`/`for`/`while`.
- **Manage Randomness**: Use explicit PRNG keys with `jax.random.PRNGKey` for reproducible randomness; split keys for multiple uses.
- **Optimize Data Types and Memory**: Default to `jnp.float32`; avoid unnecessary copies and use static shapes where possible.
- **Immutability First**: Never mutate arrays in-place; JAX arrays are immutable—create new ones for updates.
- **Profile and Test**: Validate inputs with shape/type checks; use `jax.debug.print` for JIT debugging; write pytest unit tests; benchmark with `jax.profiler`.
- **Style and Docs**: Follow PEP 8, snake_case for vars/functions, UPPER_CASE constants; add NumPy-style docstrings; keep functions small and modular.
- **Parallelism**: Use `jax.pmap` for multi-device execution when scaling across GPUs/TPUs.

### Examples
**JIT-Compiled Loss Function**:
```python
import jax
import jax.numpy as jnp
from jax import jit, grad

def loss_fn(params, x, y):
    predictions = jnp.dot(x, params)  # Pure, vectorized
    return jnp.mean((predictions - y) ** 2)

jitted_loss = jit(loss_fn)
loss_grad = grad(jitted_loss, argnums=0)  # Gradient w.r.t. params
```

**Vectorized Over Batches with vmap**:
```python
def single_predict(params, x):
    return jnp.dot(x, params)

batched_predict = jax.vmap(single_predict, in_axes=(None, 0))  # Vectorize over x batch

# Usage: predictions = batched_predict(weights, batch_x)
```

**Random Key Splitting for Reproducibility**:
```python
key = jax.random.PRNGKey(0)
key1, key2 = jax.random.split(key, 2)
weights = jax.random.normal(key1, (100, 10))
bias = jax.random.normal(key2, (10,))
```

**Scan for RNN-Like Loop**:
```python
def step_fn(carry, x):
    hidden, _ = carry
    hidden = jnp.tanh(jnp.dot(x, hidden))
    return (hidden, None), hidden

final_hidden, outputs = jax.lax.scan(step_fn, (init_hidden, None), inputs)