1. Home
  2. /
  3. Jax
  4. /
  5. Neural Network Training Example

Neural Network Training Example - Jax Typing CST Test

Loading…

Neural Network Training Example — Jax Code

Training a small neural network using JAX and gradient descent.

import jax.numpy as jnp
from jax import grad

# Data
X = jnp.array([[1.0],[2.0],[3.0],[4.0]])
y = jnp.array([2.0,4.0,6.0,8.0])

# Parameters
w = 0.0
b = 0.0

# Prediction
def predict(w, b, X):
    return w * X + b

# Loss
def loss(w, b):
    y_pred = predict(w, b, X)
    return jnp.mean((y - y_pred)**2)

grad_loss = grad(loss, argnums=(0,1))

# Gradient descent
for _ in range(1000):
    dw, db = grad_loss(w, b)
    w -= 0.01 * dw
    b -= 0.01 * db

print('Learned w,b:', w, b)

Jax Language Guide

JAX is an open-source Python library for high-performance numerical computing, combining NumPy-like API with automatic differentiation (autograd), GPU/TPU acceleration, and composable function transformations for machine learning and scientific computing.

Primary Use Cases

  • ▸High-performance machine learning and deep learning model development
  • ▸Gradient-based optimization and automatic differentiation
  • ▸Physics simulations and scientific computing requiring differentiable functions
  • ▸Research in reinforcement learning and generative models
  • ▸GPU/TPU accelerated numerical computing at scale

Notable Features

  • ▸NumPy-compatible API with hardware acceleration
  • ▸Automatic differentiation with `grad`
  • ▸Just-in-time compilation (`jit`) for optimized performance
  • ▸Vectorization (`vmap`) and parallelization (`pmap`) of functions
  • ▸Composable transformations for advanced research pipelines

Origin & Creator

JAX was developed by researchers at Google Research starting in 2018, building on Autograd and XLA (Accelerated Linear Algebra) to enable high-performance, differentiable programming.

Industrial Note

JAX is widely used in cutting-edge machine learning research, physics simulations, reinforcement learning, and differentiable programming, where composable gradients and hardware acceleration are essential.

Quick Explain

  • ▸JAX provides a NumPy-compatible API with hardware acceleration on CPU, GPU, and TPU.
  • ▸It offers automatic differentiation for gradients of arbitrary Python functions using `grad`, `vmap`, `jit`, and `pmap`.
  • ▸JAX enables composable transformations like vectorization, parallelization, and just-in-time compilation, making it ideal for research and large-scale ML experiments.

Core Features

  • ▸Autograd for forward- and reverse-mode differentiation
  • ▸JIT compilation for CPU/GPU/TPU performance
  • ▸Vectorized operations over batches with `vmap`
  • ▸Parallel computation across devices with `pmap`
  • ▸Random number generation with functional, reproducible API

Learning Path

  • ▸Learn Python and NumPy fundamentals
  • ▸Understand functional programming principles
  • ▸Practice autograd and `grad` on simple functions
  • ▸Experiment with `jit`, `vmap`, and `pmap`
  • ▸Build research or ML pipelines using Flax/Optax/JAX

Practical Examples

  • ▸Compute gradient of a scalar function with `grad`
  • ▸Train a simple neural network using JAX arrays and `grad`
  • ▸Vectorize loss computation over a batch with `vmap`
  • ▸JIT-compile a physics simulation for GPU execution
  • ▸Parallelize reinforcement learning environment rollout across multiple GPUs using `pmap`

Comparisons

  • ▸JAX vs NumPy: JAX adds autograd, JIT, vmap, pmap, GPU/TPU support
  • ▸JAX vs TensorFlow: JAX is functional, research-focused, flexible; TensorFlow has higher-level APIs and production tools
  • ▸JAX vs PyTorch: JAX uses functional programming and composable transformations; PyTorch is imperative and popular for production
  • ▸JAX vs Numpy+Autograd: JAX is faster, supports hardware acceleration and composable transforms
  • ▸JAX vs MATLAB: JAX is Python-first, differentiable, and GPU/TPU compatible

Strengths

  • ▸Extremely fast and hardware-optimized for large computations
  • ▸Highly composable functional transformations
  • ▸Seamless integration with NumPy and SciPy
  • ▸Strong support for research in ML and differentiable programming
  • ▸Works efficiently on TPUs and multi-GPU clusters

Limitations

  • ▸Steep learning curve for beginners in functional programming style
  • ▸Limited ecosystem compared to TensorFlow or PyTorch for high-level models
  • ▸Debugging JIT-compiled code can be tricky
  • ▸Some Python libraries are incompatible with JAX’s functional transformations
  • ▸Primarily research-focused; fewer production deployment utilities

When NOT to Use

  • ▸Quick prototyping for simple computations on CPU
  • ▸Projects needing extensive high-level ML framework support
  • ▸Non-research applications with no need for gradient computation
  • ▸Legacy codebases not compatible with functional transformations
  • ▸Very small scripts where JIT or GPU acceleration overhead outweighs benefits

Cheat Sheet

  • ▸Array = JAX array (like NumPy)
  • ▸grad(f) = derivative of function f
  • ▸jit(f) = compiled function for speed
  • ▸vmap(f) = vectorized map over batches
  • ▸pmap(f) = parallel map across devices

FAQ

  • ▸Is JAX free?
  • ▸Yes - open-source under Apache 2.0 license.
  • ▸Which devices are supported?
  • ▸CPU, GPU, and TPU (via XLA backend).
  • ▸Can JAX compute gradients automatically?
  • ▸Yes - using `grad` for scalar or vector functions.
  • ▸Is JAX suitable for deep learning?
  • ▸Yes - often used with Flax or Haiku for neural networks.
  • ▸Can JAX scale to multiple devices?
  • ▸Yes - `pmap` allows parallelization across GPUs/TPUs.

30-Day Skill Plan

  • ▸Week 1: NumPy-like computations and arrays
  • ▸Week 2: Automatic differentiation with `grad`
  • ▸Week 3: JIT compilation and benchmarking
  • ▸Week 4: Vectorization with `vmap` and parallelization with `pmap`
  • ▸Week 5: Full ML pipelines with Flax/Optax and TPU/GPU acceleration

Final Summary

  • ▸JAX is a high-performance Python library for numerical computing with autograd, JIT, and hardware acceleration.
  • ▸Enables functional, composable, and differentiable programming.
  • ▸Ideal for ML research, scientific computing, and TPU/GPU-accelerated pipelines.
  • ▸Integrates seamlessly with Optax, Flax, and Haiku for deep learning.
  • ▸Supports vectorization, parallelization, and scalable numerical computation.

Project Structure

  • ▸scripts/ - JAX model and function scripts
  • ▸datasets/ - input data for training or simulations
  • ▸notebooks/ - exploratory computation and experimentation
  • ▸models/ - saved parameters or checkpoints
  • ▸logs/ - performance metrics and experiment tracking

Monetization

  • ▸Research consulting using JAX pipelines
  • ▸Scientific simulations for enterprise clients
  • ▸ML model development and optimization
  • ▸High-performance computing services
  • ▸Training workshops and tutorials

Productivity Tips

  • ▸Use JIT to accelerate heavy computations
  • ▸Vectorize functions instead of Python loops
  • ▸Keep functions pure for composability
  • ▸Cache and reuse compiled functions
  • ▸Leverage multi-device parallelism with `pmap`

Basic Concepts

  • ▸Array: JAX’s primary data structure, similar to NumPy arrays
  • ▸grad: computes derivatives of functions automatically
  • ▸jit: compiles Python functions for optimized execution
  • ▸vmap: vectorizes functions over batch dimensions
  • ▸pmap: parallelizes functions across multiple devices

Official Docs

  • ▸https://jax.readthedocs.io/
  • ▸https://github.com/google/jax

More Jax Typing Exercises

JAX Simple Linear Regression ExampleJAX Logistic Regression ExampleJAX Neural Network Forward Pass ExampleJAX Mean Squared Error ExampleJAX Gradient Computation ExampleJAX Vectorized Operations ExampleJAX JIT Compilation ExampleJAX Softmax Classification ExampleJAX Convolution Example

Practice Other Languages

CReactPythonC++RustTypeScriptKotlinPHPJavaC#RubyMqlCqlN1qlCypher