Fix mx.sort vjp to transpose the permutation - #3700

Proposed changes

Closes #3699.

Sort::vjp delegated to Sort::jvp, but sorting applies a permutation P to the input and a permutation matrix is not symmetric. The jvp gathers the tangents along the argsort indices (P v); the vjp must scatter the cotangents back to the original positions (P^T w). Delegating to the jvp routed the cotangents the wrong way, giving incorrect gradients for any input whose sort permutation is not an involution. This affects CPU and GPU (the rule is in primitives.cpp), and there were no autograd tests for sort, so it went unnoticed.

The fix computes the vjp with put_along_axis on the argsort indices, mirroring the existing, correct Partition::vjp.

import mlx.core as mx
x = mx.array([3.0, 1.0, 2.0, 5.0, 4.0])
cotan = mx.array([10.0, 20.0, 30.0, 40.0, 50.0])
print(mx.vjp(lambda a: mx.sort(a), (x,), (cotan,))[1][0])
# before: [20, 30, 10, 50, 40]   after: [30, 10, 20, 50, 40] (matches finite differences)

Adds test_sort_grad covering the exact gradient and the jvp/vjp adjoint identity across axes.

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)