Introduction to JAX

JAX is a numerical computing python library developed by Google designed for high-performance machine learning research and scientific computing. It provides functionalities similar to NumPy but includes automatic differentiation, Just-In-Time (JIT) compilation, and efficient execution on GPUs and TPUs.

These features allow for optimization of mathematical computations, making JAX useful in deep learning, probabilistic modeling, and physics simulations.

Understanding JIT Compilation

Just-In-Time (JIT) compilation is a technique that improves runtime performance. In JAX, jax.jit compiles Python functions into optimized machine code using the XLA compiler. When a function is JIT-compiled, it runs significantly faster because redundant Python overhead is removed.

This optimization is useful for machine learning models, which involve repeated computations on large datasets.

Gradients and Automatic Differentiation

A gradient represents the rate of change of a function with respect to its variables. In optimization, gradients help find the direction in which a function decreases or increases.

JAX provides automatic differentiation through jax.grad, which computes gradients without requiring manual differentiation. This feature is essential in training machine learning models where gradient-based optimization algorithms like gradient descent adjust parameters iteratively to minimize error.

Implementing Linear Regression with JAX

Linear regression is a statistical method that models the relationship between an independent variable x and a dependent variable y using a linear function:

where w is the slope and b is the intercept. The goal is to find values of w and b that minimize the difference between predicted and actual values.

Generating Synthetic Data

import jax.numpy as jnp
import jax
import matplotlib.pyplot as plt

# Generate synthetic data
key = jax.random.PRNGKey(0)
x = jax.random.uniform(key, (100,), minval=0, maxval=10)
true_w, true_b = 2.5, 1.0  # True slope and intercept
y = true_w * x + true_b + jax.random.normal(key, (100,))  # Adding some noise

A dataset is created where y is computed as a linear function of x, with some noise added for realism. The true values of w and b are set to 2.5 and 1.0, respectively.

Defining the Model and Loss Function

# Define the linear model
def model(w, b, x):
    return w * x + b

# Define the loss function (Mean Squared Error)
def loss_fn(w, b, x, y):
    y_pred = model(w, b, x)
    return jnp.mean((y_pred - y) ** 2)

The model function computes predictions using w and b. The loss function calculates the mean squared error (MSE), measuring how well predictions match the actual values.

Computing Gradients and Updating Parameters

# Compute gradients
grad_fn = jax.grad(loss_fn, argnums=(0, 1))  # Compute gradients for w and b

# Gradient Descent Update
def update(w, b, x, y, lr=0.01):
    dw, db = grad_fn(w, b, x, y)
    w = w - lr * dw
    b = b - lr * db
    return w, b

JAX computes gradients automatically using jax.grad, which eliminates the need for manual differentiation. The update function applies gradient descent, adjusting w and b to reduce loss.

Training the Model

# Training loop
epochs = 100
learning_rate = 0.05
for epoch in range(epochs):
    w, b = update(w, b, x, y, learning_rate)

The model is trained for 100 iterations. At each step, the update function modifies w and b to minimize loss.

Visualizing Results

# Plot results
plt.scatter(x, y, label="Data")
plt.plot(x, model(w, b, x), color='red', label="Fitted Line")
plt.legend()
plt.xlabel("x")
plt.ylabel("y")
plt.title("Linear Regression with JAX")
plt.show()

print(f"Learned parameters: w = {w:.3f}, b = {b:.3f}")
print(f"True parameters: w = {true_w}, b = {true_b}")

A scatter plot of the data and the fitted regression line is generated. The trained values of w and b are compared to the true values.

FULL CODE:

import jax.numpy as jnp
import jax
import matplotlib.pyplot as plt

# Generate synthetic data
key = jax.random.PRNGKey(0)
x = jax.random.uniform(key, (100,), minval=0, maxval=10)
true_w, true_b = 2.5, 1.0  # True slope and intercept
y = true_w * x + true_b + jax.random.normal(key, (100,))  # Adding some noise

# Initialize parameters
w = jnp.array(0.0)
b = jnp.array(0.0)

# Define the linear model
def model(w, b, x):
    return w * x + b

# Define the loss function (Mean Squared Error)
def loss_fn(w, b, x, y):
    y_pred = model(w, b, x)
    return jnp.mean((y_pred - y) ** 2)

# Compute gradients
grad_fn = jax.grad(loss_fn, argnums=(0, 1))  # Compute gradients for w and b

# Gradient Descent Update
def update(w, b, x, y, lr=0.01):
    dw, db = grad_fn(w, b, x, y)
    w = w - lr * dw
    b = b - lr * db
    return w, b

# Training loop
epochs = 100
learning_rate = 0.05
for epoch in range(epochs):
    w, b = update(w, b, x, y, learning_rate)

# Plot results
plt.scatter(x, y, label="Data")
plt.plot(x, model(w, b, x), color='red', label="Fitted Line")
plt.legend()
plt.xlabel("x")
plt.ylabel("y")
plt.title("Linear Regression with JAX")
plt.show()

print(f"Learned parameters: w = {w:.3f}, b = {b:.3f}")
print(f"True parameters: w = {true_w}, b = {true_b}")

Why JAX is Efficient for This Task

JAX’s key benefits in this linear regression implementation include:

• Automatic Differentiation: jax.grad computes gradients efficiently, reducing errors from manual differentiation.
• Optimized Execution: Functions can be JIT-compiled for improved speed.
• GPU and TPU Acceleration: JAX can leverage specialized hardware for enhanced performance.
• Functional API: JAX encourages immutable operations, improving reproducibility and efficiency.

Other Use Cases for JAX

Beyond linear regression, JAX is used in:

• Deep Learning: Many frameworks like Flax and Haiku use JAX for training neural networks.
• Bayesian Inference: JAX powers probabilistic modeling libraries like NumPyro.
• Physics Simulations: JAX accelerates computations in fluid dynamics, quantum mechanics, and differential equations.
• Reinforcement Learning: JAX enables efficient policy gradient methods for training agents.

Thank you for reading this article. I hope you found it helpful and informative. If you have any questions, or if you would like to suggest new Python code examples or topics for future tutorials, please feel free to reach out. Your feedback and suggestions are always welcome!

Happy coding!
Py-Core.com Python Programming

You can also find this article at Medium.com