pub fn linear_last_dim(layer: &Linear, xs: &Tensor) -> Result<Tensor>Expand description
Apply a [Linear] across the last dimension of a tensor of any rank.
candle’s Linear::forward matmuls directly, which caps the input at rank 3
— beyond that it fails with shape mismatch in matmul. LigandMPNN’s ligand
graph is rank 5 ((B, L, M, M, C): one M x M edge matrix per residue), so
y_edges and every projection inside DecLayerJ needs this
(ferritin-100.11).
Flattens the leading axes, applies the layer, and restores the shape with the new feature width.