Study Plan

DiT Patchify and adaLN-Zero Block

Diffusion ModelsTransformersConditioning

Hard

Cut each image into patch tokens and apply one adaLN-Zero transformer block, returning both the tokens and the block output.

\begin{aligned} z &= \operatorname{patches}(x)\, W_e + b_e + P, \qquad [\,s_{\text{shift}}, s_{\text{scale}}, g\,] = \operatorname{SiLU}(c)\, W_a + b_a \\ \text{out} &= z + g \odot \Big( \big( \operatorname{LN}(z) \odot (1 + s_{\text{scale}}) + s_{\text{shift}} \big) W_m + b_m \Big) \end{aligned}

With p = patch_size and x of shape (B, C, height, width), the image is split into a (height/p) by (width/p) grid of p by p patches. Token n = r (width/p) + q holds the patch in grid row r and column q, flattened in (C, p, p) order, so entry k p p + i p + j is x[b, k, r p + i, q p + j]. The weights come from params with D the token width and E the conditioning width.

$W_e$ is embed_weight (C p p, D), $b_e$ is embed_bias (D), $P$ is pos_embed (N, D).
$c$ is conditioning, $W_a$ is ada_weight (E, 3D) and $b_a$ is ada_bias (3D). The modulation of shape (B, 3D) splits along its last axis into shift, scale and gate in that order, each (B, D) and broadcast over the N tokens.
$W_m$ is mlp_weight (D, D) and $b_m$ is mlp_bias (D).

LN normalizes each token over its D features with the population variance and eps = 1e-6, with no learned weight or bias. Return a dict with keys "tokens" ( $z$ ) and "output", each a tensor of shape (B, N, D) with the dtype of x.

Examples

Input: x = [[[[1.0,0.0,2.0,-1.0],[0.0,1.0,1.0,0.0],[-1.0,0.5,0.0,0.0],[0.5,-0.5,1.0,2.0]]]], conditioning = [[0.5,-1.0]], params = {"embed_weight":[[0.5,-0.2,0.1],[0.3,0.4,-0.1],[-0.2,0.1,0.6],[0.1,0.2,0.3]],"embed_bias":[0.0,0.1,-0.1],"pos_embed":[[0.1,0.0,0.0],[0.0,0.1,0.0],[0.0,0.0,0.1],[0.1,0.1,0.1]],"ada_weight":[[0.2,-0.1,0.3,0.5,0.0,-0.4,0.3,0.2,0.1],[0.1,0.4,-0.2,0.0,0.3,0.2,-0.5,0.1,0.2]],"ada_bias":[0.0,0.1,0.0,0.2,-0.1,0.0,0.5,0.5,0.5],"mlp_weight":[[0.5,0.0,-0.3],[0.2,0.4,0.1],[-0.1,0.3,0.6]],"mlp_bias":[0.1,0.0,-0.1]}, patch_size = 2

Output: {"tokens":[[[0.7,0.1,0.3],[0.5,-0.5,0.8],[-0.5,0.45,0.0],[0.1,0.7,1.2]]],"output":[[[1.317031,-0.107485,-0.076611],[0.554508,-0.600033,0.877892],[-0.904346,0.682028,0.28646],[-0.514448,0.884871,1.714136]]]}

Explanation: The top-left 2x2 patch flattens to (1, 0, 0, 1), so its token is rows 0 and 3 of embed_weight summed, plus embed_bias and pos_embed[0], giving (0.7, 0.1, 0.3). The output adds the gated, modulated projection of each normalized token.

Input: x = [[[[1.0,2.0,3.0,4.0],[5.0,6.0,7.0,8.0]],[[-1.0,0.0,1.0,0.0],[0.0,1.0,0.0,-1.0]]]], conditioning = [[1.0,0.0]], params = {"embed_weight":[[0.1,0.0],[0.0,0.1],[0.1,0.1],[-0.1,0.0],[0.5,-0.5],[0.0,0.2],[0.3,0.0],[-0.2,0.4]],"embed_bias":[0.0,0.0],"pos_embed":[[0.0,0.5],[0.5,0.0]],"ada_weight":[[0.1,0.2,0.3,-0.3,1.0,0.5],[0.4,0.4,0.4,0.4,0.4,0.4]],"ada_bias":[0.0,0.0,0.0,0.0,0.0,0.0],"mlp_weight":[[1.0,0.0],[0.0,1.0]],"mlp_bias":[0.0,0.0]}, patch_size = 2

Output: {"tokens":[[[-0.7,2.1],[1.4,0.2]]],"output":[[[-1.537948,2.438807],[2.344836,-0.031917]]]}

Input: x = [[[[0.5,-0.5],[1.0,2.0]]],[[[0.0,1.0],[-1.0,0.0]]]], conditioning = [[1.0,2.0],[-1.0,0.5]], params = {"embed_weight":[[1.0,-1.0]],"embed_bias":[0.0,0.5],"pos_embed":[[0.0,0.0],[0.1,0.1],[0.2,0.2],[0.3,0.3]],"ada_weight":[[0.0,0.0,0.0,0.0,0.0,0.0],[0.0,0.0,0.0,0.0,0.0,0.0]],"ada_bias":[0.0,0.0,0.0,0.0,0.0,0.0],"mlp_weight":[[0.7,0.2],[-0.3,0.9]],"mlp_bias":[0.1,-0.1]}, patch_size = 1

Output: {"tokens":[[[0.5,0.0],[-0.4,1.1],[1.2,-0.3],[2.3,-1.2]],[[0.0,0.5],[1.1,-0.4],[-0.8,1.7],[0.3,0.8]]],"output":[[[0.5,0.0],[-0.4,1.1],[1.2,-0.3],[2.3,-1.2]],[[0.0,0.5],[1.1,-0.4],[-0.8,1.7],[0.3,0.8]]]}

Hint 1

A reshape to (B, C, height/p, p, width/p, p) followed by a permute separates the grid position from the position inside a patch.

Hint 2

F.conv2d with stride=p applies the patch embedding once the weight is reshaped to (D, C, p, p); check the embed_weight layout in the description before reshaping.

Hint 3

The function torch.nn.functional.layer_norm takes an eps argument and works without weight and bias.

Hint 4

tensor.chunk(3, dim=-1) splits the modulation vector into three equal parts.

Requirements

Use PyTorch tensor operations; torch.nn.functional is allowed.
Leave the inputs, including every tensor in params, unchanged.

Constraints

x is a float tensor of shape (B, C, height, width) with B between 1 and 3, C between 1 and 3, and height and width between 2 and 8, both divisible by patch_size.
patch_size is an integer between 1 and 4, and N = (height / patch_size) (width / patch_size) is between 1 and 16.
conditioning has shape (B, E) with E between 1 and 5, and the token width D is between 2 and 8.
params holds exactly the keys embed_weight (C p p, D), embed_bias (D), pos_embed (N, D), ada_weight (E, 3D), ada_bias (3D), mlp_weight (D, D) and mlp_bias (D).
Every tensor has the same dtype as x, which is float32 or float64. All values are finite, and ada_weight and ada_bias may be all zeros.

Try Similar Problems

Toy Diffusion Training StepHard Toy DDPM SamplerHard Frechet Distance Between Feature SetsMedium Tiny U-Net Forward PassHard DDIM StepHard Karras Euler SamplerHard Gaussian Mixture ScoreMedium Cosine Noise ScheduleMedium Vision Transformer Patch EmbeddingHard Transformer Encoder BlockHard Implement PagedAttention Block AllocationHard Scaled Dot-Product AttentionHard