affine_grid

→Tensor

affine_grid(theta: Tensor, size: list[int] | tuple[int, ...], align_corners: bool | None = None)

source

Generate a sampling grid from a batch of affine transform matrices.

Builds the flow field needed to apply an affine transform via grid_sample. Together they form a Spatial Transformer Network (STN, Jaderberg et al. 2015), where theta is typically the output of a localisation network.

For each output pixel $(x_{\mathrm{out}}, y_{\mathrm{out}})$ the corresponding source coordinate is computed as:

\begin{pmatrix} x_{\mathrm{src}} \\ y_{\mathrm{src}} \end{pmatrix} = \theta \, \begin{pmatrix} x_{\mathrm{out}} \\ y_{\mathrm{out}} \\ 1 \end{pmatrix}

Parameters

thetaTensor

Batch of affine matrices of shape (N, 2, 3) (2-D) or (N, 3, 4) (3-D). Each row encodes one element of the affine transform applied to homogeneous output coordinates.

sizesequence of int

Target output shape passed to grid_sample afterwards — (N, C, H, W) for 2-D, (N, C, D, H, W) for 3-D.

align_cornersbool= None

Must match the align_corners argument later passed to grid_sample. Controls whether -1/+1 refer to the centre of the corner pixels (True) or to the outer edge (False, default).

Returns

Tensor

Sampling grid of shape (N, H, W, 2) (or (N, D, H, W, 3) for 3-D) in normalised $[-1, 1]$ coordinates, ready to be fed to grid_sample.

Notes

affine_grid does not look at the source image — only at the requested output size — so the resulting grid is reusable across any input of matching spatial dims and any channel count.

Examples

>>> import lucid
>>> from lucid.nn.functional import affine_grid
>>> theta = lucid.tensor([[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]])  # identity
>>> g = affine_grid(theta, (1, 3, 4, 4), align_corners=False)
>>> g.shape
(1, 4, 4, 2)