KeOps: a symbolic tensor that never materialises the matrix it represents

The problem is that your program is memory bound, not compute bound

The KeOps documentation starts from an observation about how mathematical software represents data, and the observation is more useful than the feature list.

Dense matrices are the default. Variables are encoded as dense numerical arrays, which is convenient and well supported, but puts a heavy load on memory. Large arrays are expensive to move around and may not even fit in RAM or GPU memory. The consequence stated in the documentation is that a majority of scientific programs are memory bound: run times for most neural networks and mathematical computations are not limited by the raw capabilities of the CPUs and CUDA cores, but by the time consuming transfers of large arrays from memory circuits to arithmetic computing units.

That is the diagnosis. If your program is spending its time on transfers rather than arithmetic, making the arithmetic units faster will not help, and the only real fix is to do less transferring.

The second half of the diagnosis is about the usual workaround. The common answer is sparse matrices, tensors with few non zero coefficients, represented as lists of indices and values. The documentation grants that this method is elegant and allows large arrays with a small memory footprint, then explains the catch: it does not stream well on GPUs. Parallel computing devices are wired for block wise memory access and struggle with lists of random indices. As a consequence, compared with dense arrays, sparse encodings only speed up computations for matrices with less than one percent non zero coefficients, and that restrictive condition keeps sparse matrices mostly useful for graph and mesh processing.

So the choice is dense and slow, or sparse and inapplicable. KeOps proposes a third option.

Symbolic matrices, where the coefficient is a formula and not a number

The key remark in the documentation is that most of the large arrays used in machine learning and applied mathematics share a common mathematical structure. Distance matrices, kernel matrices, point cloud convolutions and attention layers can all be described as symbolic tensors: given two collections of vectors, their coefficient at position (i,j) is given by a mathematical formula evaluated on the data samples.

The distinction from sparse is precise. These objects are not sparse in the traditional sense, because most coefficients are non zero, but they can still be described efficiently using a formula and relatively small data arrays.

KeOps therefore stores the formula and the data, not the matrix. What the user holds is a symbolic tensor whose shape is nominally the full matrix, and the library is responsible for never building it.

The implementation note explains how that is possible. The library takes advantage of the structure of CUDA registers to bypass costly memory transfers between arithmetic and memory circuits. A reduction over a symbolic matrix is turned into a fused kernel that reads the input vectors once, computes the formula, and accumulates the result, so the only memory traffic is the input data and the output vector. That is the same fusion strategy a good hand written CUDA kernel would use, which is why the speed-up claim in the documentation is stated as x10 to x100 against a PyTorch GPU baseline on geometric applications.

The important caveat is that the comparison is against a dense PyTorch implementation. Where the baseline is already efficient, or where the operation set is small, the gap narrows.

The LazyTensor API, and why the reshape looks strange

The Python example in the documentation is worth reading closely, because the odd shapes are the whole point.

The first step creates two arrays on the GPU, one with a million rows and one with two million, three columns each.

import torch
M, N, D = 1000000, 2000000, 3
x = torch.randn(M, D, requires_grad=True).cuda()
y = torch.randn(N, D).cuda()

Then the dense tensors are turned into KeOps symbolic variables with virtual dimensions at positions zero and one, for the i and j indices.

from pykeops.torch import LazyTensor
x_i = LazyTensor(x.view(M, 1, D))
y_j = LazyTensor(y.view(1, N, D))

The reshape to (M, 1, D) and (1, N, D) looks backwards, and it is. A symbolic tensor has two virtual dimensions, and the data is placed along a third. That is what lets the library know which index is i and which is j, and it is why the operations that follow are written against dimension 2 for the feature axis.

From there the computation is ordinary looking.

D_ij = ((x_i - y_j)**2).sum(dim=2)
K_ij = (- D_ij).exp()

a_i = K_ij.sum(dim=1)

The first line forms a symbolic matrix of squared distances of shape (1e6, 2e6, 1), the second a symbolic Gaussian kernel matrix of the same shape, and the third reduces along the j dimension to produce a real tensor of shape (1e6, 1). The documentation notes that the reduction is computed using a CUDA scheme with a linear memory footprint, and that in the kernel density estimation example it outperforms standard PyTorch implementations by two orders of magnitude.

Automatic differentiation works through the same abstraction.

g_x = torch.autograd.grad((a_i ** 2).sum(), [x])

The gradient of a function of a never-materialised matrix is computed by differentiating the fused kernel, not by storing the matrix and differentiating over it. That is the feature that makes the library usable inside a PyTorch training loop rather than only as a one-off computation.

Python, Matlab and R, and three language bindings in one repository

The library combines efficient C++ routines with an automatic differentiation engine and can be used with Python (NumPy, PyTorch), Matlab and R.

The repository layout shows how that is organised. There is a keopscore/ directory holding the core C++ implementation, a pykeops/ directory for the Python package, an rkeops/ directory for R, a keops.Rproj file for the R project, and a benchmarks/ directory. The docs live in doc/ with a doc-requirements.txt for the documentation dependencies, and there is a pydoc.sh and a pybuild.sh for the build steps.

The Python package is published on PyPI as pykeops, and the badges in the readme show the PyPI version and download counts, while the CRAN badges for rkeops are commented out in the current readme, which suggests the R distribution is not currently advertised there even though the rkeops directory is present.

The website is at kernel-operations.io, and the readme points there for documentation, contribution guidelines and tutorials, with the documentation tree including API pages, an auto tutorials index and an auto benchmarks index. Those three separate things, API reference, tutorials and benchmarks, are the minimum for a library whose performance is the reason to use it.

One point about the project structure deserves emphasis. The readme explains that the Docker image also includes the libraries hosted on the same website, such as GeomLoss, and argues that this is convenient and has negligible impact on the final image size, since Cuda and PyTorch weigh about 5Gb anyway and there is little point in maintaining separate images that differ by a handful of Python files. That is a pragmatic packaging decision, and it also means the container is not a minimal KeOps image.

A compiled library means the install is a build, not a wheel download

The Dockerfile is the clearest evidence of what it takes to build KeOps, and it is a heavy build.

The image is built by./docker-images.sh, and the Dockerfile takes a series of build arguments rather than pinning versions directly: BASE_IMAGE, PYTHON_VERSION, NUMPY_VERSION, KEOPS_VERSION which the comment calls the most important parameter, GEOMLOSS_VERSION, CUDA_VERSION, PYTORCH_VERSION, and then PYTEST_VERSION, HYPOTHESIS_VERSION, JAXTYPING_VERSION, BEARTYPE_VERSION and BLACK_VERSION for the test and formatting toolchain.

The comments explain why the versions are arguments. The Python version argument exists to test support across Python versions. The PyTorch scatter dependency is noted as possibly lagging PyTorch releases by a few days, with a pointer to the pytorch_scatter project for compatibility information, and the ARG for its version is commented out. The testing stack is listed explicitly: KeOps relies on PyTest, Hypothesis, Beartype and Jaxtyping for unit tests, and Black for code formatting.

That last detail is worth pausing on. Combined with a compiled C++ core, this is a library that treats correctness of shapes and types as a first class concern, which makes sense when the whole abstraction rests on a formula being evaluated on the right indices.

The practical consequence for an adopter is that there is a compilation step. The Dockerfile installs build-essential, and a source installation needs a working C++ toolchain and a compatible CUDA setup. A pure pip install of a prebuilt wheel is the easy path where one exists for your platform, and the CI matrix that varies the Python version argument exists to keep that coverage honest.

Where symbolic tensors are the wrong abstraction

The honest limitations follow from the abstraction itself, and they are worth stating before anyone tries it.

The first is that the operation set is restricted. A symbolic tensor supports a limited family of algebraic combinations, and above all the reductions, operations such as .sum(), .logsumexp() and .argmin(), that bring it back to a real tensor. The documentation says this explicitly: you come back to vanilla PyTorch tensors or NumPy arrays using reduction operations on one of the two symbolic dimensions. If your computation needs an arbitrary operation on the matrix, or a matrix product against a dense array, or an operation that is not a reduction, the symbolic form may not carry it. That is a real constraint on the design, not an oversight.

The second is that the win depends on the baseline. The x10 to x100 speed-up is stated against a PyTorch GPU baseline on geometric applications, from kernel methods to geometric deep learning. If your workload is small enough to fit comfortably in device memory, the dense path is fine and the compilation overhead is pure loss. The benefit grows with the size of the problem, and it peaks exactly where the dense representation fails.

The third is that the abstraction leaks. You write expressions like ((x_i - y_j)**2).sum(dim=2) with a hard coded dimension index, and that index encodes the convention that the feature axis is dimension two. Code written this way is coupled to the layout of the symbolic tensor, and it does not look like ordinary tensor code. Readability is the price, and it is a real price in a codebase other people have to maintain.

The fourth is the build weight. A compiled core, a CUDA toolchain, and a dependency on PyTorch versions being in step means KeOps is a heavier thing to deploy than a pure Python package, and the Docker image the project maintains is accordingly large.

Editorial conclusion

KeOps fits a research or production team whose kernel, distance or point cloud computation is memory bound, where the dense baseline runs out of device memory or spends most of its time moving data, and who need gradients through the same computation. It does not fit a codebase that needs an arbitrary operation on the matrix, because the symbolic abstraction only supports a restricted set of algebraic combinations and reductions, and it does not fit a team unwilling to accept a compilation step in the workflow. Check three things before you commit. Check that your expression can be written in the supported form, since the symbolic tensor is defined by a formula and the permitted operations are limited to what the compiler can handle. Check your build toolchain, because the library compiles C++ kernels at install time and a source install needs a working compiler rather than a pure wheel. And check which backend you are on, since the speed-up claim in the project documentation is stated against a PyTorch GPU baseline, so a CPU-only deployment sees a different comparison. The last push was on 2026-07-17, the newest release is v2.3 from 2025-04-18, and the license is MIT.

Frequently asked questions

Official sources

1. getkeops/keops on GitHub
2. License: MIT
3. Project website
4. README
5. Releases