Module linear 

Functions

burn_linear_bias

y = x @ w^T + b. x: [lead, d_in], w: [d_out, d_in], b: [d_out].

burn_linear_nd

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.

burn_linear_nobias

y = x @ w^T (no bias).