Implement one step of SGD with momentum.

Given parameters w, gradient grad, velocity v, learning rate lr, and momentum coefficient mu, compute:

$v_{new} = \mu v + grad$

$w_{new} = w - lr \cdot v_{new}$

Write a function momentum_step(w, grad, v, lr, mu) that returns the tuple (w_new, v_new) as PyTorch tensors. Do not modify the inputs in place; return new tensors.