Fix axis param in nll_loss - #3651

Proposed changes

nll_loss ignores its axis argument. The gather happens along axis but the targets get expanded at the last dim and the result squeezed at the last dim, so anything other than the default axis either fails with a confusing squeeze error, or returns wrong values with no error at all if the trailing dim happens to be 1:

import mlx.core as mx
import mlx.nn as nn

logp = mx.log(mx.softmax(mx.random.normal((3, 4)), axis=0))
nn.losses.nll_loss(logp, mx.array([0, 2, 1, 2]), axis=0)
# ValueError: [squeeze] Cannot squeeze axis -1 with size 4 which is not equal to 1.

cross_entropy had the same problem and was fixed in #2641, but nll_loss kept the old pattern. This applies the same fix and adds a test along the lines of the one from that PR. Checked the result against numpy take_along_axis for axis 0, a middle axis on a 3d input, and the default axis.

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)