AI Frameworks

Machine learning frameworks provide the foundational software layer between model definitions and hardware execution. They abstract the complexity of tensor operations, automatic differentiation, GPU memory management, and distributed computation, allowing researchers and engineers to focus on model design rather than low-level implementation details.

The Big Three: PyTorch, TensorFlow, and JAX

The framework landscape has consolidated around three major players: PyTorch, TensorFlow, and JAX. Each embodies different design philosophies with distinct trade-offs. PyTorch prioritizes developer experience and flexibility. TensorFlow emphasizes production deployment and ecosystem completeness. JAX focuses on functional programming, composability, and high-performance computing.

Framework choice has profound implications for the entire ML stack. It influences debugging workflows, deployment options, hardware compatibility, community support, and the availability of pre-trained models and libraries. Organizations must consider not just the current feature set but the long-term trajectory and community momentum of each framework.

Specialized Frameworks and Runtimes

Beyond the big three, specialized frameworks serve specific domains. ONNX Runtime provides cross-framework inference optimization. TensorRT optimizes for NVIDIA GPU deployment. Core ML and TFLite Micro target mobile and embedded devices respectively. Understanding this ecosystem is essential for making informed infrastructure decisions.

• ONNX Runtime — Cross-framework inference optimization via a standardized model format.
• TensorRT — NVIDIA's high-performance inference optimizer for GPU deployment.
• Core ML — Apple's framework for on-device ML on iOS, macOS, and Apple Silicon.
• TFLite Micro — TensorFlow's runtime for microcontrollers and embedded devices.
• OpenVINO — Intel's toolkit for optimizing inference on Intel hardware.

The "framework wars" have settled into ecosystem niches — PyTorch for research, TensorFlow for production, JAX for HPCDeeper InsightRather than one framework winning outright, the ML community has converged on a division of labor. PyTorch dominates research and rapid prototyping thanks to its Pythonic API and dynamic graphs. TensorFlow retains its grip on production deployment with its unmatched breadth of serving targets (mobile, web, embedded, cloud). JAX has carved out the high-performance computing niche, favored by Google DeepMind and teams training the largest foundation models on TPU pods. Understanding these niches helps teams pick the right tool for their specific context. Click to collapse

PyTorch has become the dominant framework in ML research due to its intuitive Pythonic interface and dynamic computational graph (eager execution). Every operation executes immediately, making debugging as simple as inserting print statements or using standard Python debuggers. This define-by-run approach matches how researchers think about models.

Autograd: Automatic Differentiation

PyTorch's autograd system automatically records operations on tensors and computes gradients through reverse-mode automatic differentiation. The computational graph is rebuilt on every forward pass, enabling dynamic architectures that vary from input to input, such as recursive neural networks or models with data-dependent control flow.

import torch

class="tok-comment"># Autograd example: compute gradients automatically
x = torch.tensor([class="tok-number">2.0, class="tok-number">3.0], requires_grad=True)
y = x ** class="tok-number">2 + class="tok-number">3 * x  class="tok-comment"># y = x^class="tok-number">2 + 3x
loss = y.sum()
loss.backward()      class="tok-comment"># Compute dy/dx
print(x.grad)         class="tok-comment"># tensor([class="tok-number">7., class="tok-number">9.])  (2x + class="tok-number">3)

From Research to Production: torch.compile

For production deployment, PyTorch offers TorchScript (a subset of Python that can be JIT compiled) and torch.export for ahead-of-time compilation. PyTorch 2.0 introduced torch.compile, which uses a graph capture and compilation approach to achieve significant speedups without requiring code changes. This bridges the gap between PyTorch's research flexibility and production performance.

The PyTorch Ecosystem

The PyTorch ecosystem includes Hugging Face Transformers for pre-trained models, PyTorch Lightning for training boilerplate reduction, and torchvision/torchaudio/torchtext for domain-specific functionality. This rich ecosystem, combined with strong community momentum, makes PyTorch the default choice for most new ML projects.

• Hugging Face Transformers — Pre-trained models and fine-tuning pipelines for NLP, vision, and multimodal tasks.
• PyTorch Lightning — Reduces training boilerplate and standardizes project structure.
• torchvision / torchaudio / torchtext — Domain-specific datasets, transforms, and model architectures.
• FSDP (Fully Sharded Data Parallel) — Native distributed training with memory-efficient parameter sharding.

TensorFlow was designed from the ground up for production ML at Google scale. Its original design centered on static computational graphs that are defined once and then executed repeatedly, enabling aggressive compiler optimizations, distributed execution, and deployment across diverse hardware targets.

TensorFlow 2.0 and Keras

TensorFlow 2.0 adopted eager execution by default, matching PyTorch's developer experience. The tf.function decorator allows selectively tracing Python functions into optimized graphs, providing a middle ground between ease of use and performance. Keras, integrated as TensorFlow's high-level API, provides a user-friendly interface for common model architectures.

import tensorflow as tf

class="tok-comment"># tf.function traces Python into an optimized graph
class="tok-decorator">@tf.function
def train_step(model, x, y):
    with tf.GradientTape() as tape:
        predictions = model(x, training=True)
        loss = loss_fn(y, predictions)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return loss

The Deployment Ecosystem

TensorFlow's strongest advantage is its comprehensive deployment ecosystem. No other framework matches the breadth of deployment targets that TensorFlow supports.

For large-scale training, TensorFlow offers robust distributed strategies through tf.distribute, supporting data parallelism, model parallelism, and parameter server architectures. Integration with TPUs provides access to Google's custom ML accelerators, which can offer significant cost and performance advantages for large training jobs.

JAX brings a functional programming paradigm to ML computing, built on top of XLA (Accelerated Linear Algebra) compiler. JAX programs are written as pure functions that transform data, which enables powerful composition of transformations like automatic differentiation, vectorization, parallelization, and JIT compilation.

Composable Transformations

JAX's core transformations are composable and orthogonal. These transformations can be applied independently or chained together, enabling patterns that would be cumbersome in other frameworks.

import jax
import jax.numpy as jnp

class="tok-comment"># Composable transformations: grad + jit + vmap
def loss_fn(params, x, y):
    pred = jnp.dot(x, params)
    return jnp.mean((pred - y) ** class="tok-number">2)

class="tok-comment"># Compose: JIT-compile the gradient of the loss
fast_grad = jax.jit(jax.grad(loss_fn))

class="tok-comment"># Compose: vectorize over a batch of different param sets
batched_grad = jax.vmap(jax.grad(loss_fn), in_axes=(class="tok-number">0, None, None))

The JAX Ecosystem

JAX's ecosystem is growing rapidly with libraries like Flax and Haiku for neural network modules, Optax for optimization, and JAX-based implementations of major model architectures. Google DeepMind has standardized on JAX for research, and it is increasingly used for large-scale training of foundation models.

• Flax — Google's neural network library for JAX with a focus on flexibility.
• Optax — Gradient processing and optimization library (SGD, Adam, etc.).
• Orbax — Checkpointing and serialization utilities for JAX models.
• Haiku — DeepMind's lightweight neural network library for JAX.

Computational graphs are the fundamental abstraction underlying all ML frameworks. They represent the sequence of mathematical operations that transform inputs into outputs, with nodes representing operations and edges representing data flow. The structure of this graph determines what optimizations are possible.

Static vs. Dynamic Graphs

Static graphs (TensorFlow 1.x, compiled modes) enable whole-program optimization: the compiler can see all operations before any execution occurs, allowing it to fuse operations, eliminate redundancy, optimize memory allocation, and plan device placement. The trade-off is reduced flexibility and harder debugging.

Dynamic graphs (PyTorch eager, TensorFlow eager) build the graph on-the-fly during execution, providing maximum flexibility for dynamic control flow and easy debugging. Modern frameworks bridge this gap through tracing-based compilation (torch.compile, tf.function) that captures dynamic graphs and compiles them for subsequent executions.

ML Compilers: The Emerging Layer

The trend toward ML compilers is reshaping the framework landscape. Projects like Apache TVM, MLIR, and Triton provide compiler infrastructure that can optimize ML workloads across diverse hardware. These compilers sit between frameworks and hardware, enabling a clean separation of concerns where framework developers focus on usability and compiler developers focus on performance.

• XLA (Accelerated Linear Algebra) — Google's ML compiler, used by JAX and TensorFlow for TPU and GPU targets.
• TorchInductor — PyTorch 2.0's default compiler backend, generating Triton kernels for GPU execution.
• Apache TVM — Open-source compiler stack for deploying ML to diverse hardware targets.
• MLIR (Multi-Level IR) — LLVM-based infrastructure for building domain-specific compilers, increasingly used for ML.
• Triton — OpenAI's language for writing custom GPU kernels in Python-like syntax.