pykeops

2.3 verified Mon Apr 27 auth: no python

PyKeOps is a Python package providing bindings to the KeOps library, which enables efficient computation of kernel operations on CPUs and GPUs with automatic differentiation and without memory overflows. It is widely used in geometry processing, shape analysis, and machine learning for large-scale distance and kernel computations. Current version is 2.3, with active development on GitHub.

pip install pykeops

Common errors

error RuntimeError: KeOps CUDA error: no kernel image is available for execution on the device ↓

cause The precompiled CUDA binaries do not support your GPU architecture (e.g., too old or too new).

fix

Reinstall pykeops with environment variable KEOPS_GPU_ARCH set to your GPU architecture, e.g., KEOPS_GPU_ARCH=sm_86 pip install pykeops --no-cache-dir

error ImportError: cannot import name 'LazyTensor' from 'pykeops' ↓

cause Importing from wrong module; LazyTensor is in pykeops.torch, not pykeops.

fix

Use: from pykeops.torch import LazyTensor

Warnings

breaking PyKeOps 2.0 introduced a completely reworked compilation engine. Old code using from pykeops.torch import generic_logsumexp or similar may not work. Update to LazyTensor API. ↓

fix Migrate to LazyTensor-based API: replace generic reductions with symbolic expressions.

gotcha LazyTensor operations must be explicitly evaluated (e.g., .eval(), .detach(), or .sum() returns a LazyTensor, not a torch.Tensor). Forgetting evaluation leads to symbolic objects in the result. ↓

fix Call .eval() on the final LazyTensor or use reduction operations that return tensors (some like .sum() still return LazyTensor).

deprecated The old API using pykeops.torch.generic_sum, generic_logsumexp, etc. is deprecated in 2.x and may be removed. Use LazyTensor or Genred instead. ↓

fix Use LazyTensor or Genred for custom formulas.

Imports

LazyTensor
wrong
```
from pykeops import LazyTensor
```
correct
```
from pykeops.torch import LazyTensor
```
LazyTensor is in the torch submodule; importing directly from pykeops fails.
genred
wrong
```
from pykeops import genred
```
correct
```
from pykeops.torch import Genred
```
Genred is the correct spelling, lowercase genred does not exist.

Quickstart

Compute pairwise Gaussian kernel sum using LazyTensors for memory-efficient reduction.

import torch
from pykeops.torch import LazyTensor

# Create sample data
x = torch.randn(1000, 3).cuda()
y = torch.randn(2000, 3).cuda()

# Define LazyTensors
X = LazyTensor(x[:, None, :])  # (1000, 1, 3)
Y = LazyTensor(y[None, :, :])  # (1, 2000, 3)

# Compute pairwise distances squared
D = ((X - Y) ** 2).sum(2)  # Symbolic reduction

# Compute softmin (log-sum-exp) with Gaussian kernel
K = (-D).exp()
result = K.sum(1)  # (1000,) actually returns a LazyTensor, evaluate with .detach()
print(result.detach())