Function burn_linear_nd 

pub fn burn_linear_nd<const D: usize>(
    x: Tensor<D>,
    w: Tensor<2>,
    b: Option<Tensor<1>>,
) -> Tensor<D>

y = x @ w^T + b for rank-D input: flatten leading dims, apply burn_linear_bias or burn_linear_nobias, restore shape with the last dim replaced by w’s output dim.