tinygrad-notes

Convolution windows and arange

All tutorials · Updated September 21, 2026 · tinygrad 8ad8f73

Convolution and sequence construction illustrate how familiar operations can be assembled from smaller Tensor operations. Their frontend definitions do not tell you the final kernel count: rewrites and scheduling can remove or fuse intermediate work.

Convolution is a reduction over windows

from tinygrad import Tensor

x = Tensor.arange(9).reshape(1, 1, 3, 3).float().realize()
weight = Tensor.ones(1, 1, 2, 2).realize()
result = x.conv2d(weight)
assert result.shape == (1, 1, 2, 2)
assert result.tolist() == [[[[8.0, 12.0], [20.0, 24.0]]]]
print(result.tolist())

The top-left output sums 0+1+3+4. Moving the window right produces 1+2+4+5. Each output reduces the input-channel and filter dimensions, while batch, output channel, and output spatial positions survive.

In OpMixin.conv2d, the general path pads the input and calls _pool in mixin/movement.py to describe windows. It reshapes and broadcasts those windows across output channels, multiplies by reshaped weights, then reduces the channel/filter dimensions.

_pool builds the window mapping with movement operations. It is not necessarily an allocation containing a separate copy of every window. Rangeification later translates those mappings into indices.

Groups partition input and output channels. Stride controls window origins; dilation spaces filter samples. These alter the indexing problem and should be checked separately when implementing or changing a kernel.

The implementation also has IMAGE and Winograd branches. The simple general-path explanation is not a claim that every configured convolution follows it.

arange begins as cumulative addition

from tinygrad import Tensor

assert Tensor.arange(2, 11, 3).tolist() == [2, 5, 8]
assert Tensor.arange(5, -1, -2).tolist() == [5, 3, 1]
assert Tensor.arange(0).tolist() == []
print("positive, negative, and empty ranges match")

The current arange calculates a length, constructs a constant tensor containing step, applies cumulative addition, adds start-step, and casts to the requested dtype. It also checks representability and handles an empty range.

This is the frontend expression. It does not mean execution performs a general parallel scan for every arange: symbolic/compiler simplification can exploit the constant construction. Inspect generated source for the actual program rather than estimating runtime cost from the Python method alone.

A useful experiment is to compare an arange used only for indexing with one explicitly realized as data. These consumers impose different storage requirements.

Original chapter by Di Zhu: historical version.