torchrolling

Pandas-style rolling and exponentially weighted statistics for PyTorch tensors, fast on GPU. Only depends on torch.

import torch
import torchrolling

x = torch.tensor([1.0, 2.0, float("nan"), 4.0, 5.0])
torchrolling.rolling(x, 3, min_periods=2).mean()
# tensor([nan, 1.5000, 1.5000, 3.0000, 4.5000])

prices = torch.randn(512, 10_000, device="cuda").cumsum(-1)  # [series, time]
returns = prices.diff(dim=-1, prepend=prices[..., :1])
r = torchrolling.rolling(returns, 60)
features = torch.stack([r.mean(), r.std(), r.skew(), r.median(), r.corr(returns.roll(1, -1))])
trend = torchrolling.ewm(prices, span=20).mean()

Why

pandas.Series.rolling works on one column on the CPU. When your series already live in a [batch, time] tensor on the GPU (features computed on the fly, inside a model, after augmentation), going through pandas means a round trip per batch. The usual torch workaround, x.unfold(-1, w, 1).mean(-1), does w times more work than needed, and for a median or quantile it copies a tensor w times bigger than x. torch has no rolling median and no exponential moving average at all.

torchrolling computes rolling sums, means, counts, variances, skew, kurtosis, covariances, correlations, minima and maxima in O(1) work per element and O(n) memory for any window size; medians and quantiles; and exponentially weighted means, variances and correlations with a parallel scan. On CUDA every statistic is a single fused Triton kernel. It works on any device, supports autograd and torch.compile, and gives the same numbers as pandas.

Install

pip install torchrolling

Python 3.10+, torch 2.0+. The CUDA kernels need Triton 3.2+, which comes with torch 2.6+; with an older torch, the same statistics run as plain torch operations.

Usage

Rolling windows

torchrolling.rolling(x, window, *, min_periods=None, center=False, dim=-1, acc_dtype=None) returns an object with these methods. Each returns a tensor shaped like x.

Call several methods on one object: counts and moments are computed once and shared.

Exponentially weighted windows

torchrolling.ewm(x, com=None, span=None, halflife=None, alpha=None, *, min_periods=0, adjust=True, ignore_na=False, dim=-1, acc_dtype=None) takes exactly one of com, span, halflife, alpha, as in pandas.

Semantics

The test suite checks everything against pandas with property-based tests (hypothesis):

• the output has the same length; windows that are not full yet give NaN unless min_periods allows them;
• NaN and ±inf are treated as missing and skipped (pandas does the same). count() is the exception, as in pandas: it counts ±inf;
• center=True centres the window the way pandas does;
• pairwise statistics (cov, corr) use only the positions where both series are valid;
• integer input becomes the default float dtype; float input keeps its dtype;
• interpolation is one of linear, lower, higher, midpoint, nearest, as in pandas.

Where results can differ from pandas:

• ewm(adjust=False) with missing values and alpha=0.5 (also com=1, span=3, halflife=1): pandas takes a different formula for that case (pandas-dev/pandas#66523); torchrolling uses the documented weights, as pandas does for every other alpha;
• ewm(...).var() and .cov() with bias=False, when the weights have decayed a lot (alpha close to 1, long runs of missing values): torchrolling computes the bias correction without cancellation, so it can differ from pandas from the 8th digit on, and gives NaN rather than 0 after a single observation;
• corr is NaN, not ±inf, where one series is constant in the window.

Precision: acc_dtype

Sums and moments are accumulated in the input's precision, but at least float32: float64 input is accumulated in float64 and matches pandas to 1e-9 or better; float32, float16 and bfloat16 input is accumulated in float32, which is accurate to about 1e-6 relative. Minima, maxima and quantiles are exact in the input dtype. To accumulate float32 data in float64, pass acc_dtype=torch.float64. On consumer and inference GPUs (T4, RTX), float64 is much slower than float32.

Gradients

Every statistic supports autograd. On CUDA, the fused Triton kernels compute the forward pass when no gradient is needed (inference, feature pipelines, torch.no_grad()); when one is, torchrolling runs the same algorithms as plain torch operations instead, which autograd differentiates. Both give the same results.

How it works

The series is padded and cut into blocks of exactly window elements. Each block is scanned forwards (prefix) and backwards (suffix) with cumsum, cummax or cummin. Any window then covers the tail of one block and the head of the next, so its value is combine(suffix[start], prefix[end]) (the van Herk / Gil-Werman algorithm). For sums this also means only 2 * window numbers are ever added together, so error does not grow with series length.

Moments (variance, skew, kurtosis, covariance) come from power sums, which cancel badly when taken raw. Each side of a window is therefore shifted by a value that lies inside that same side: a head starts at its block's start, so it contains the block's first valid value, and a tail contains its block's last one. The shifted sums are as well conditioned as the window itself, a constant window gives exact zeros (so its variance is exactly 0, as in pandas), and the two sides are merged with the pairwise formulas of Chan, Golub and LeVeque and of Pébay.

Quantiles cannot be split into halves; the torch code sorts every window, in chunks so that memory stays bounded.

On CUDA, each of these is a single fused Triton kernel. For rolling statistics, a program loads a run of whole blocks twice, once as is and once shifted by one block and one element, so that the tail and the head of every window sit at the same position of two register tiles; both scans and the merge happen in registers, in one pass over memory. For quantiles with windows above 64, a program sorts the segment its outputs need once, with each value's position packed into the sort key, then walks the sorted segment once, counting for every window how many of its own values it has passed: O(window) work per output instead of O(window log² window). Smaller windows sort each window directly.

Exponentially weighted statistics follow pandas' update rule mean_t = (1 - s_t) * mean_{t-1} + s_t * x_t. The weights s_t depend only on where values are missing, so they are computed up front, and the mean, the weighted covariance and the bias correction become affine recurrences y_t = a_t * y_{t-1} + b_t. These are composed in parallel in blocks of 64 steps (Hillis-Steele), which is stable because every a_t lies in [0, 1]. On CUDA, one program per series scans it chunk by chunk and carries the state across chunks.

Benchmarks

Milliseconds, best of 3 runs, lower is better; the fastest in each row is bold. float32 data, measured with bench/bench.py. The GPU columns ran on a Tesla T4 (Kaggle, torch 2.11, Triton 3.6, via notebooks/kaggle_gpu.ipynb); pandas and polars ran twice, on that Kaggle machine's CPU and on an Apple M3 Max (16 cores; pandas 3.0, polars 1.44). Each library gets its data where it wants it (tensors on the GPU, a DataFrame for pandas and polars), and only the computation is timed. For EWM, window is the span. All raw numbers, including float64 accumulation, are in bench/results/.

1,000 series × 10,000 points

10,000 series × 10,000 points (GPU libraries only)

torch unfold is the usual plain-torch workaround, for example x.unfold(-1, w, 1).mean(-1): it has no missing-value handling and no min_periods, and for median and corr it copies a tensor window times the size of the input. On small windows without missing values it is the fastest way to get mean, std or max; everywhere else torchrolling is. "—" means the library has no such statistic.

Reading these numbers honestly. Much of the gap to pandas and polars is a GPU against a CPU, and it depends on the CPU: on the M3 Max, polars runs 5–7 times faster than on Kaggle's CPU. One comparison goes the other way. For rolling median with windows of 200 and 1000, polars on the M3 Max takes 39 and 37 ms, against torchrolling's 44 and 81 ms on the T4 (a 2018 inference GPU). Those polars times leave out what it takes to get data that already lives on the GPU into polars: copying it to the CPU and the result back (40 MB each way for this table) and building the DataFrame. So where the data lives decides: if it is already on the GPU, sending it to the CPU rarely pays off; if it is on the CPU, weigh the copy to the GPU, and staying on the CPU with pandas or polars may well be the better choice. cuDF is measured at its weakest shape, thousands of short columns, which it processes one by one; on a few long series it would be much closer. The first call of each statistic, which compiles its Triton kernel (a few seconds, then cached on disk), is not timed either.

The kernels' launch configurations were tuned on a T4. Results are the same on every GPU; only speed can differ. On a GPU or system where a kernel cannot run, torchrolling warns once and computes the same statistics with plain torch operations.

License

MIT