Pair-tile VJP #

Layer.pairTile (the distogram stem): two residue blocks Xi, Xj : Mat L D — the host's features, constants of the step — and two weights W, Wj : Mat D C; the pair map is y i j c = (Xi · W) i c + (Xj · Wj) j c. It is linear in each weight, so each weight's VJP is the adjoint of a linear map: dW = Xiᵀ · (Σ_j dy) and dWj = Xjᵀ · (Σ_i dy) — the cotangent summed over the axis the weight was broadcast along, then the dense weight rule. These are the two stablehlo.reduce + dot_general pairs the pairTile backward emits (%ptb_du / %d_W, %ptb_dv / %d_Wj). The layer takes no input gradient: its input is the feature block.

Both witnesses are pdiv_of_affine with the other block's term as the constant, the shape of pdiv_dense_W; the only work is the row-major (i, (j, c)) index of the flat pair map.

Host pair planes (Layer.pairTile's pairIn) #

With K host pair planes the layer's output is the tile with the planes P : Vec (L·(L·K)) appended — constants of the step, like the feature blocks. In the flat layout the plane block follows the tile block (finSumFinEquiv): the output is tileW … v t at inl t and P q at inr q. Each weight's Jacobian is the tile's own on the first block and zero on the second, so each VJP is the tile's backward on the cotangent's tile block (tileBlock) — the stablehlo.slice of the first C channels the pairTile backward emits when pairIn > 0.