[nastya236]

Observation:

For a linear layer we do matrix multiplication in the following way:
x @ W.T, x \in [M, K], W \in [N, K], K is fast moving dim, x is the input, W is the weight.
In vjp we compute the gradient of the loss with respect of the weights as:
dW = x^T @ dY (dY incoming gradient \in [M, N]) -> [K, N] gradient, N is fast moving. Then Transpose::vjp just changing the strides, shape is [N, K], N is still fast moving.
=> there is a mismatch between the layout of dW (row major) and layout of W that it column major
=> we use suboptimal strided kernel to do optimizer step W-= dW (simplified).

Fix:

This PR proposes to write gradients in a matching layout to the second operand: if B was produced using Transpose primitive we compute dW = dY^T @ x, else the opposite.

Results:

For 7B pretraining on 1 B200 with batch size 1, context 8192 this change results in 17500 tok/s -> 19400 tok/s throughput change.