AI frameworks: PyTorch, TensorFlow and JAX

In short: AI frameworks take over tensor calculation, automatic differentiation and use of the graphics card. PyTorch is the standard in research today, TensorFlow and JAX are the main alternatives.

In more detail: You describe the model as a sequence of operations, the framework generates the gradients when needed and runs them on GPU or TPU.

In detail

• PyTorch (Meta): dynamic graphs, Python-like, huge ecosystem (Hugging Face, Lightning). torch.compile speeds it up.
• TensorFlow/Keras (Google): especially for production, mobile devices (TensorFlow Lite) and web (TensorFlow.js). Keras is the simple interface.
• JAX (Google): functional style, jit, grad, vmap. Popular for research on TPUs.
• Hugging Face Transformers: ready-made models and training code, see Hugging Face.
• Inference: ONNX Runtime, TensorRT, llama.cpp (GGUF and llama.cpp), vLLM, MLX (Apple), WebGPU in the browser.
• Minimal PyTorch example:

import torch
model = torch.nn.Linear(3, 1)
opt = torch.optim.Adam(model.parameters(), lr=1e-2)
for x, y in data:
    loss = ((model(x) - y) ** 2).mean()
    opt.zero_grad(); loss.backward(); opt.step()

• Hardware drivers: CUDA (NVIDIA), ROCm (AMD), Metal (Apple).

See also: Forward and backward pass in detail, Why GPUs for AI?