tinygrad-notes

Shapes, movement operations, and indexing

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

Current tinygrad does not use the old ShapeTracker/View classes. Shape transformations live in the UOp graph. During scheduling, rangeification turns logical coordinates into expressions that select values from source tensors. This chapter replaces the old ShapeTracker tutorial at the same URL.

A transpose changes the coordinate mapping

For a complete composition rather than isolated operations, follow the transpose/pad/reduction trace. It derives 3*k + i - 4, checks its validity domain, and shows the actual rangeified UOps. Its regression oracle also exercises reversal, asymmetric padding, and broadcasting over 15 shapes.

Consider row-major data with two rows and three columns:

from tinygrad import Tensor
from tinygrad.uop.ops import Ops

x = Tensor([[0, 1, 2], [3, 4, 5]]).realize()
y = x.permute(1, 0)
assert y.shape == (3, 2)
assert y.uop.op is Ops.PERMUTE
assert y.tolist() == [[0, 3], [1, 4], [2, 5]]
print(y.tolist())

In the original tensor, coordinate (row, column) selects offset 3*row + column. After transposition, output coordinate (i, j) refers to source coordinate (j, i), hence offset 3*j + i.

The transpose describes a mapping. It does not imply an immediate copy, nor guarantee that a later consumer needs no copy. Scheduling, required layout, and the consumer determine materialization.

Follow the implementation of a mapping

apply_movement_op in schedule/indexing.py propagates coordinates backward through movement operations. Its arguments include the source shape, the movement argument, and the coordinates used to index the result.

from tinygrad.uop.ops import UOp, Ops
from tinygrad.schedule.indexing import apply_movement_op

i = UOp.variable("i", 0, 2)
j = UOp.variable("j", 0, 1)
source = apply_movement_op(Ops.PERMUTE, (2, 3), (1, 0), (i, j))
assert source == (j, i)
offset = source[0] * 3 + source[1]
assert offset.sym_infer({"i": 2, "j": 1}) == 5
print(offset.render())

This example calls the actual indexing helper, rather than reconstructing an obsolete View object. The printed expression may reorder terms; its meaning is the tested coordinate mapping.

Movement Backward coordinate mapping
Permute Reorder result coordinates using the inverse permutation.
Shrink Add the slice’s starting offset to the corresponding coordinate.
Flip Replace coordinate r by size - 1 - r on flipped axes.
Reshape Flatten result coordinates, then divide/remainder by source dimensions.
Expand Remove introduced broadcast coordinates at the internal representation’s boundary.
Pad Subtract the padding offset and attach validity conditions.

Public Tensor broadcasting also involves shape normalization. Do not assume that the internal EXPAND argument is interchangeable with the public Tensor shape argument; inspect the graph and movement implementation.

Reshape is flattening followed by unflattening

Reshaping a contiguous (2, 3) tensor to (3, 2) maps output (i, j) to flat index q = 2*i + j. The source coordinates are (q // 3, q % 3).

from tinygrad.uop.ops import UOp, Ops
from tinygrad.schedule.indexing import apply_movement_op

i = UOp.variable("i", 0, 2)
j = UOp.variable("j", 0, 1)
row, col = apply_movement_op(Ops.RESHAPE, (2, 3), (3, 2), (i, j))
for a in range(3):
  for b in range(2):
    bindings = {"i": a, "j": b}
    assert row.sym_infer(bindings) * 3 + col.sym_infer(bindings) == 2*a+b
print("all six coordinates match")

_apply_reshape builds this arithmetic and simplifies it with symbolic rewrite rules. The source explicitly identifies that simplification as the replacement for reshape view merging. Composed transformations can cancel or simplify before the final program is rendered.

Padding adds a validity condition

For an input of length three padded by one on each side, output coordinate r is valid when 1 <= r < 4; valid elements use input coordinate r - 1. Outside that region the result is zero.

from tinygrad import Tensor

x = Tensor([10, 20, 30])
assert x.pad(((1, 1),)).tolist() == [0, 10, 20, 30, 0]
assert x.shrink(((1, 3),)).flip(0).tolist() == [30, 20]
print("padding, slicing, and reversal match")

convert_pad_to_where_to_keep_behavior_local uses validity to select input values or zero. A correct lowering must preserve this distinction; replacing a padded access with an unconditional out-of-bounds load is not equivalent.

From coordinates to kernels

run_rangeify determines realization boundaries and propagates ranges through consumers and producers. It rewrites reductions to refer to explicit ranges, expresses padding with selection, and removes movement nodes once their effects have been incorporated into indexing.

get_kernel_graph then simplifies the graph, converts staged values into stores, splits kernels, and removes views from the kernel graph. This is why a Tensor-level transpose need not survive as a transpose instruction.

Continue with dimension merging and range splitting and upcasting.

Original chapter by Di Zhu: historical version.