JIT Compilation & Optimization

Overview

Just-In-Time (JIT) compilation makes Python fast enough for serious computation:

• PyPy: Alternative Python runtime with tracing JIT
• Numba: JIT compile numerical Python to machine code (100-1000x speedup!)
• Cython: Write C-like Python that compiles to C
• Type specialization: JIT infers types and specializes code
• Guard-based optimization: Assumes types stay constant

-

Numba: JIT for Numerical Code

Basic Numba Usage

from numba import jit
import time
import numpy as np

# Pure Python (slow)
def sum_loop_python(arr):
 total = 0
 for x in arr:
 total += x
 return total

# Numba JIT (fast!)
@jit
def sum_loop_numba(arr):
 total = 0
 for x in arr:
 total += x
 return total

# Benchmark
arr = np.arange(1000000)

start = time.time()
result_python = sum_loop_python(arr)
python_time = time.time() - start

start = time.time()
result_numba = sum_loop_numba(arr)
numba_time = time.time() - start

print(f"Python: {python_time:.3f}s")
print(f"Numba: {numba_time:.3f}s")
print(f"Speedup: {python_time / numba_time:.0f}x")
# Typical

Numba Type Specification

from numba import jit, float64, int64
import numpy as np

# Specify input/output types for optimization
@jit(float64(float64[:], float64[:]))
def dot_product(a, b):
 """Compute dot product."""
 total = 0.0
 for i in range(len(a)):
 total += a[i] * b[i]
 return total

a = np.array([1.0, 2.0, 3.0])
b = np.array([4.0, 5.0, 6.0])
result = dot_product(a, b) # 32.0

Numba CUDA for GPU Acceleration

from numba import cuda
import numpy as np

# Compile to GPU
@cuda.jit
def gpu_add(out, a, b):
 """Add two arrays on GPU."""
 idx = cuda.grid(1) # Get thread index
 if idx < out.size:
 out[idx] = a[idx] + b[idx]

# Use it
a = np.array([1, 2, 3, 4, 5], dtype=np.float32)
b = np.array([6, 7, 8, 9, 10], dtype=np.float32)
out = np.empty_like(a)

# Call with block/thread configuration
gpu_add[1, 32](out, a, b)
print(out) # [7. 9. 11. 13. 15.]

---

PyPy: Alternative Runtime with JIT

Why PyPy is Fast

# PyPy uses tracing JIT:
# 1. Interprets code initially
# 2. Detects hot loops
# 3. JIT compiles hot paths to machine code
# 4. Specializes for observed types

# Example where PyPy excels

def fibonacci(n):
 """Classic recursive computation."""
 if n <= 1:
 return n
 return fibonacci(n-1) + fibonacci(n-2)

# CPython
# PyPy
# Why? JIT optimizes the recursive calls

result = fibonacci(35)
print(result)

Running with PyPy

# Install PyPy (if not already)
# On macOS

# Run Python script with PyPy
pypy3 script.py

# Check if using PyPy
python -c "import sys; print(sys.implementation.name)" # 'cpython'
# With PyPy:
pypy3 -c "import sys; print(sys.implementation.name)" # 'pypy'

---

How JIT Works: Type Specialization

Type Guards

def add_numbers(a, b):
 """This function can work with different types."""
 return a + b

# CPython (no JIT)
# Calls
# Overhead

# Numba JIT
# First call
# Compiles
# Subsequent calls

# Type change (guards fail):
# add_numbers("hello", "world") → Falls back to Python

Tracing JIT (PyPy)

# PyPy traces hot code paths

def hot_loop():
 """This will be traced and JIT compiled."""
 total = 0
 for i in range(1000000):
 total += i # Hot path (executed 1M times)
 return total

# PyPy:
# 1. Initially interprets (slow)
# 2. Detects hot loop (loop counter)
# 3. Records trace
# 4. Compiles trace to machine code
# 5. Replaces hot path with compiled code
# 6. Result

-

Optimization Strategies

Write JIT-Friendly Code

from numba import jit
import time
import numpy as np

# SLOW
def slow_version():
 arr = np.arange(1000000)
 result = []
 for x in arr:
 if x % 2 == 0:
 result.append(x * x)
 return sum(result)

# FAST
@jit
def fast_version():
 arr = np.arange(1000000)
 result = 0
 for x in arr:
 if x % 2 == 0:
 result += x * x
 return result

# JIT guidelines:
# 1. Simple operations (arithmetic, comparisons)
# 2. Loops over arrays
# 3. Type consistency (no type changes in loop)
# 4. Avoid Python object creation

Avoid JIT Overhead

from numba import jit
import time
import numpy as np

# Overhead
@jit
def jit_function(n):
 return sum(range(n))

# First call (includes compilation)
start = time.time()
result1 = jit_function(10000)
first_call = time.time() - start

# Subsequent calls (uses compiled code)
start = time.time()
for _ in range(1000):
 result = jit_function(10000)
subsequent_calls = time.time() - start

print(f"First call: {first_call*1000:.1f}ms (includes compilation)")
print(f"1000 calls: {subsequent_calls*1000:.1f}ms ({subsequent_calls/1000*1000:.2f}ms each)")

-

Real-World ML Example: Custom Loss Function

Without JIT

import torch
import time

def custom_loss_python(predictions, targets):
 """Custom loss function in pure Python."""
 loss = 0
 for i in range(len(predictions)):
 error = predictions[i] - targets[i]
 if error > 0:
 loss += error ** 2
 else:
 loss += error ** 0.5
 return loss

# Slow for large batches
predictions = torch.randn(10000)
targets = torch.randn(10000)

start = time.time()
for _ in range(100):
 loss = custom_loss_python(predictions.numpy(), targets.numpy())
python_time = time.time() - start
print(f"Python loop: {python_time:.2f}s")

With Numba JIT

from numba import jit
import numpy as np
import torch
import time

@jit
def custom_loss_numba(predictions, targets):
 """Same function but JIT compiled."""
 loss = 0
 for i in range(len(predictions)):
 error = predictions[i] - targets[i]
 if error > 0:
 loss += error ** 2
 else:
 loss += error ** 0.5
 return loss

# Fast JIT version
predictions_np = np.random.randn(10000).astype(np.float64)
targets_np = np.random.randn(10000).astype(np.float64)

start = time.time()
for _ in range(100):
 loss = custom_loss_numba(predictions_np, targets_np)
numba_time = time.time() - start
print(f"Numba JIT: {numba_time:.2f}s")
print(f"Speedup: {python_time / numba_time:.0f}x")

---

Comparison: JIT Methods

Method	Speedup	Use Case	Ease
Numba @jit	100-1000x	Numerical loops	Easy
Numba CUDA	1000-10000x	GPU acceleration	Medium
PyPy	10-100x	Long-running processes	Hard to integrate
Cython	10-100x	Mixed Python/C	Medium
NumPy/PyTorch	10-1000x	Vector operations	Easy

-

When to Use JIT

Use Numba When

CPU-bound loops with numerical operations Need 100-1000x speedup Can't use NumPy/PyTorch Loop runs many times

Don't Use Numba When

Code is already fast (NumPy, PyTorch) Lots of Python object creation Complex logic with many branches Single-execution code

Use PyPy When

Long-running processes (servers, simulations) Pure Python code (no C extensions) Want automatic optimization Can't modify code

---

Summary: JIT Performance

Execution speed (relative):

CPython: 1x (baseline)
PyPy: 10-100x (tracing JIT)
Numba: 100-1000x (type specialization)
NumPy: 10-100x (vectorization + C)
PyTorch GPU: 1000-10000x (parallel GPUs)
Numba CUDA: 1000-10000x (specialized GPU code)

-

Related Topics

• 01 Python Bytecode Fundamentals - What gets JIT compiled
• [04 Profiling & Performance Analysis](/05-py3/09-bytecode-and-execution/(04-profiling-performance-analysis/) - Measuring JIT impact
• 00 Readme - Memory during JIT execution
• 00 Readme - Production optimization