Encyclopedia · 187 concepts

MLOps & Infrastructure · beginner · concept 136 of 187

ML Frameworks (PyTorch, TensorFlow, JAX)

A framework supplies three things a hand-written model would need: automatic differentiation, meaning gradients computed from a recorded graph of operations rather than derived by hand; tensor ops dispatched to vendor GPU kernels; and primitives for splitting a model across devices. PyTorch carries most research code, JAX is the transformation-based option behind much of Google DeepMind's work, TensorFlow persists in older production and mobile stacks. The common error is treating the choice as a performance decision. Speed comes from kernels, compilers and memory layout, and portability off NVIDIA hardware is where frameworks actually differ.

Key terms

automatic differentiationcomputation graphGPU kernelsXLA compilationeager vs graph execution

Where you meet it in the real world

Training and serving nearly every production model: LLM pretraining runs, vision models in autonomous driving stacks, recommendation systems, on-device speech and camera features via mobile runtimes