Fix vmap of partition/argpartition dropping the kth argument

[Adityaj0]

Fixes #4113.

Proposed changes

Partition::vmap and ArgPartition::vmap called the two-argument overload

partition(inputs[0], axis_ + axis_left, stream())

which is the one that partitions the flattened array, with axis_ + axis_left bound to kth. kth_ was dropped entirely, so under vmap these ops returned a 1-D array partitioned around the wrong element. Sort::vmap immediately below is correct only because sort has no kth parameter.

topk is built on partition, so it inherited the bug and failed outright once the traced slice was applied to the flattened result.

Passing kth_ explicitly selects the (a, kth, axis, stream) overload.

Before, on main:

vmap(partition)    -> (24,)   expected (2, 3, 4)
vmap(argpartition) -> (24,)   expected (2, 3, 4)
vmap(topk)         -> ValueError: [slice] Invalid number of indices or strides
                                 for array with dimension 1.
vmap(topk, in_axes=1) -> IndexError: SmallVector out of range.

After:

vmap(partition)    -> (2, 3, 4), kth element matches the loop reference
vmap(argpartition) -> (2, 3, 4)
vmap(topk)         -> (2, 3, 2)
vmap(topk, in_axes=1) -> (2, 3, 2)

Tests

Added test_vmap_partition and test_vmap_topk to python/tests/test_vmap.py; there was previously no vmap coverage for these ops. They sweep every in_axes, inner axis and every valid kth / k, comparing against the equivalent loop.

Because partition only pins the kth element and leaves the two sides in an arbitrary order, the tests compare the sorted result against the sorted input and check the kth element against the reference, rather than asserting element-wise equality that the API does not promise. argpartition is checked by gathering with the returned indices.

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

Verified with a CPU-only build (-DMLX_BUILD_METAL=OFF); test_vmap, test_ops and test_autograd pass.