The Python stack for machine learning

Machine learning in practice means Python, and Python for machine learning means its libraries: a few dozen packages that between them load the data, shape it, train the model, serve it and keep track of what happened. Knowing them is most of the job. This post goes through them by use case: for each job you are likely to have, which library to reach for, what it was built for, and the sign that it is time to move to the next one. No language syntax; this is about the libraries and when they fit.

One thing to know first

Every library here is a compiled engine (C, C++, CUDA or Rust) with a Python face. Your Python code describes the work; the engine does it. That is why Python’s own slowness doesn’t matter, and it gives the one rule the whole post rests on: let the library do the work, never loop over rows yourself. When something is slow, the fix is almost never a faster library; it is finding the place where your code stepped out of the engine.

The layer that matters most is the first one, because it decides whether the others can hand data to each other without copying. More on that at the end.

Job 1: I have a table and I need to shape it

pandasA dataset that fits in memory, headed for scikit-learn

Whythe native input of every ML library; the whole tutorial corpus; the richest API for messy data

Move on whenit fits but takes minutes, or you are writing the same group-by twice

PolarsJoins, filters and aggregations on gigabytes

Whymulti-threaded by default; its lazy API plans the whole query and reads only the rows and columns it needs; a streaming engine for data bigger than memory

Move on whenyou need a scikit-learn-specific feature pandas has and Polars doesn't (convert at the end, not the start)

DuckDBA pile of Parquet or CSV files, or colleagues who think in SQL

WhySQL straight over files and over pandas or Polars frames, no server, in-process

Move on whenthe logic stops being query-shaped

Spark or DaskData that does not fit one machine

Whythe same dataframe idea, sharded across a cluster

Move on whenit fits one machine after all (most of the time it does)

Most of this layer was built by people who were annoyed with what existed. pandas started inside a hedge fund, AQR, in 2008, because its author wanted to do finance in Python and nothing handled a dated table properly; it was open-sourced the next year (pandas history). Eight years on, the same author, Wes McKinney, wrote a post titled “Apache Arrow and the ‘10 Things I Hate About pandas’”, admitting that a 5 GB file could take 20 GB of memory in his own library and that moving a table between two processes was “extremely expensive”; Arrow was his answer, started under the Apache Foundation so no vendor could own it (the post). Polars and DuckDB are, in a sense, what you build once Arrow exists: Polars from scratch in Rust to use every core, DuckDB by a database research group in the Netherlands who borrowed SQLite’s idea of “no server, just a library” and pointed it at analytics (Why DuckDB).

In practice you do not pick one. The pattern that works is DuckDB or Polars for the heavy shaping, pandas for the last mile into the model, with the handoff between them free because all three share the Arrow memory format. Mixing them costs nothing; looping over rows in any of them costs everything.

Two things in pandas worth knowing without the version history: assignment now copies rather than aliasing (so the old chained-assignment habit df[a][b] = x silently stopped working), and strings are fast only if PyArrow is installed. Install PyArrow.

Job 2: I need a model on tabular data

scikit-learnA baseline, a pipeline, a cross-validated search, an interpretable model

Whythe fit / predict / transform / Pipeline interface that every other library copies; the best documentation in the field

Move on whenthe data is tabular and accuracy is what you are paid for

XGBoost, LightGBM, CatBoostThe best accuracy on a table of features

Whygradient-boosted trees win tabular problems; all three plug into scikit-learn pipelines unchanged

Move on whenyou have images, text, audio or sequences, where trees have nothing to grip

CatBoostA table with many categorical columns

Whyhandles categories natively without you encoding them

LightGBM or XGBoostVery large tables, GPU, or a cluster

Whyboth train distributed and on GPU; LightGBM is the leaner of the two

scikit-learn began as a Google Summer of Code project in 2007 by David Cournapeau, and became what it is when four researchers at INRIA in France took it over in 2010 and made the first public release (about). That history shows in its character: academic care about correctness, documentation written to teach, and a refusal to chase every new model. XGBoost came from the other direction: Tianqi Chen, then a PhD student, wrote it to win machine learning competitions, and the 2016 paper with Carlos Guestrin (arXiv) describes the engineering that made a twenty-year-old idea, gradient boosting, fast enough to dominate them. LightGBM (Microsoft) and CatBoost (Yandex) are the two big companies’ answers to the same problem.

The thing to learn deeply here is not an algorithm but the estimator API. Once fit, predict, Pipeline and GridSearchCV are second nature, every tabular library is familiar, including ones that do not exist yet, because they will copy the same interface. The random forests and boosting posts show the mechanics on a real market dataset.

And the decision people get wrong most: a table of features wants a boosted tree, not a neural network. Reach for deep learning when the input is unstructured, not when the problem is important.

Job 3: I need a neural network

PyTorchAnything, by default

Whythe framework the open-model ecosystem is built on: Hugging Face models, LLM serving, most papers; eager and debuggable, compiled afterwards with torch.compile

Move on whenyou are on TPUs or want the compiler to see the whole program

JAXResearch code, TPUs, large sharded training, functional style

Whyjit, grad and vmap are transformations you apply to pure functions; the same code runs on CPU, GPU and TPU

Move on whenthe team writes stateful, in-place code and fights the functional style

Keras 3One codebase that must run on more than one framework

Whythe same model code runs on JAX, PyTorch or TensorFlow; saved models carry no backend

Move on whenyou need framework-specific features

TensorFlow and TF LiteAn existing TensorFlow estate, or inference on phones and devices

Whymaintained, and still the edge-deployment path

Move on whenstarting anything new

PyTorch versus JAX is a style choice more than a capability one. PyTorch feels like Python: objects with state, mutation, printing a tensor mid-training. JAX feels like maths: pure functions, no mutation, and in return the compiler can fuse, batch and shard the whole thing. Teams of software engineers usually settle into PyTorch; teams of researchers at scale often settle into JAX. Both are fine answers. TensorFlow as a first choice is no longer one.

Job 4: I need a model somebody else trained

TransformersA language, vision or audio model from the Hub, fine-tuned or as is

Whythe client for the largest model library there is; one from_pretrained call; PyTorch underneath

Move on whenyou need to serve it to many users at once

vLLMServing an LLM to concurrent users

Whypaged attention and continuous batching let one GPU serve far more requests than a request loop; speaks the OpenAI API, so clients need no change

Move on whenthe model is small enough that a plain HTTP wrapper is fine

ONNX RuntimeRunning a trained model inside a Java service, a browser, a phone, or on an NPU

Whyexport once, run anywhere, no Python runtime on the host; tuned for latency and binary size

Move on whenthe model uses operations the exporter does not understand (then you serve from Python)

llama.cpp and OllamaLocal models on your own hardware, no serving stack

Whyquantised models on CPU and consumer GPUs, one binary

Move on whenyou need throughput, not convenience

vLLM came out of UC Berkeley in June 2023, from a group who measured that existing serving systems “waste 60% – 80% of memory due to fragmentation and over-reservation” in the attention cache, and borrowed the fix from operating systems: page it, like virtual memory (vLLM announcement). The same pattern runs through this whole stack: the gain came from a systems idea, not a new model.

Two ideas in the serving layer are worth reading about even if you never run a GPU: paged attention (treating the attention cache like virtual memory, so it fragments less and more requests fit) and continuous batching (admitting new requests into a running batch instead of waiting for the batch to finish). They are why a single card now serves what used to need a rack.

Job 5: I need to keep track of what I did

MLflowWhich run produced this model, with what data and parameters

Whyevery run becomes one addressable object: parameters, metrics, artefacts, the model itself; a registry for promoting versions; tracing for LLM applications

uvAn environment that reproduces, and scripts that run anywhere

Whyone tool for environments, lockfiles and Python versions; a script can declare its dependencies in a comment at the top and run with one command

bothNumbers that survive a retrain

Whya result you cannot rerun is an anecdote; the scripts behind this site run on uv for that reason

uv is the newest tool here and the one that changed habits fastest. It comes from Astral, the people behind the Ruff linter, who announced it in February 2024 with the stated aim of a “Cargo for Python”: one tool, fast, that makes environments boring (announcement). Python packaging had been a running joke for fifteen years. Within about a year of uv it stopped being one.

For a model that a risk team or a regulator will ask about, the registry matters more than the last point of accuracy. “Model 14, trained on this data, with these parameters, approved by this person” is what a review needs, and MLflow is the cheapest way to have it.

Job 6: I need it to run faster

In order, because each step is cheaper than the next:

1. Vectorise. Replace the Python loop with the NumPy, pandas or Polars operation that does the same thing in the engine. This fixes most cases.
2. Change the engine. pandas to Polars or DuckDB for data; a boosted tree instead of a neural network for tables. Same code shape, faster execution.
3. Compile. torch.compile for PyTorch models; Numba for a numeric Python function you cannot vectorise; Cython when you need C and want to stay close to Python.
4. Parallelise. Processes, not threads, for Python-level work; the data engines are already multi-threaded underneath. Free-threaded Python exists now and the libraries are catching up, but for production this year the answer is still processes and vectorisation.
5. Move to the GPU. Only after the above, and only for work that is matrix-shaped. CuPy gives you NumPy on the GPU; scikit-learn and SciPy can already run parts of their work on PyTorch or CuPy arrays through the array API standard.

Most teams that think they need Rust needed step 1.

The seams: why these libraries can be mixed at all

Three agreements under the hood make the stack composable, and knowing them is what lets you mix libraries without paying for it:

Arrow is why DuckDB can query a Polars frame with no conversion. DLPack is why a PyTorch tensor can be handed to JAX without leaving the GPU. The array API is why the same scikit-learn code can run on a CUDA tensor. When you design a pipeline, “which library” matters less than “does the handoff copy”, and if both sides speak the same format, it does not.

The purple arrows are free. The one copy left is at the NumPy boundary into the ML libraries, and the array API standard is the project to remove it.

What is worth learning, in order

1. The scikit-learn estimator API, cold. Everything imitates it.
2. One dataframe library deeply, the other two by contract: when a handoff copies, when it doesn’t, what lazy evaluation buys.
3. Why a model that runs eagerly fails to compile. That is where production time goes.
4. The serving layer as a system: batching, the attention cache, quantisation, the OpenAI-compatible API as the integration contract.
5. Reproducibility as tooling: lockfiles, pinned versions for fast-moving libraries, a registry.

Summary

References

• Documentation of the libraries, by job:
  • tables: NumPy, pandas, Polars, DuckDB, Spark, Dask
  • tabular models: scikit-learn, XGBoost, LightGBM, CatBoost
  • neural networks: PyTorch, JAX, Keras 3, TensorFlow
  • pre-trained models and serving: Transformers, vLLM, ONNX Runtime, llama.cpp, Ollama
  • lifecycle and speed: MLflow, uv, Numba, Cython, CuPy
  • the seams: Apache Arrow, DLPack, Python array API standard
• Where the origin stories come from:
  • About scikit-learn and the history of pandas
  • Wes McKinney, “Apache Arrow and the ‘10 Things I Hate About pandas’” (2017)
  • Why DuckDB
  • Tianqi Chen and Carlos Guestrin, “XGBoost: A Scalable Tree Boosting System” (2016)
  • the vLLM announcement (2023) and the uv announcement (2024)