Fix fft vmap and jvp for transforms over a subset of axes

[kapellirohith]

Problem

The forward mode derivative of an fft that does not cover every axis of the
array is silently wrong. Since a batch axis is never transformed, that is the
ordinary batched case.

Minimal reproducer (M3 Pro, macOS 26.x, main @ e78d894):

import mlx.core as mx

mx.random.seed(0)
x = mx.random.normal((4, 5)) + 0j
t = mx.random.normal((4, 5)) + 0j
f = lambda a: mx.fft.fft(a, axis=0)

# the fft is linear, so the tangent is exactly f(t)
tangent = mx.jvp(f, [x], [t])[1][0]
print(mx.max(mx.abs(tangent - f(t))))            # before: 8.1547 ; after: 0.0
print(mx.max(mx.abs(tangent - mx.fft.fftn(t))))  # before: 0.0    ; after: 8.1547

The second line is the positive identification: before this change the tangent
is not merely wrong, it is exactly fftn over every axis, to 0.0.

A 1-D input, where the single axis is all the axes: 0.0 before and after.
mx.grad of the same function: correct before and after.
mx.fft.fftn(a, axes=(0, 1)) on a 3-D array: 16.038 before, 0.0 after.
Only forward mode, and only when the transformed axes are a strict subset.

Mechanism

FFT::jvp called the no-axes fftn, rfftn, ifftn and irfftn overloads.
Those resolve to fft_impl(a, real, inverse, norm, s), which builds
std::vector<int> axes(a.ndim()) and std::iotas it, so the tangent was
transformed over every axis regardless of axes_.

For rfft(axis=0) on a (4, 6) real input the shape itself comes back
(4, 4) instead of (3, 6): rfftn over both axes halves the last axis,
rfft(axis=0) halves the first.

FFT::vjp immediately above already threads axes_ into all four branches,
which is why reverse mode is correct and this survived.

Separately, FFT::vmap applied the real transform size change to every
transformed axis, while fft_impl resizes only valid_axes.back(). vmap of
rfft2 on a (3, 8, 8) array returned (3, 5, 5) instead of (3, 8, 5), and
vmap of irfft2 returned an output larger than the input supports. Note that
valid_axes is in the order the caller passed, not sorted, so the axis that
changes size is the last one passed; the fix keys off fft_axes.back() and
preserves that.

Both date to the initial commit d1f8627 (2023-11-29).

Why existing tests missed it

test_fft_grads covers mx.grad only, and there was no vmap or jvp coverage
for the fft anywhere in the Python or C++ suites.

Known limitation, not fixed here

vmap of irfftn with an explicitly requested odd output length still returns
2 * (n - 1):

base = mx.random.normal((3, 4)) + 0j
mx.fft.irfft(base[0], n=7).shape          # (7,)
mx.vmap(lambda a: mx.fft.irfft(a, n=7))(base).shape   # (3, 6)

This is a genuine collision, not an oversight in the rule. The primitive stores
only (axes_, inverse_, real_), and n = 6 and n = 7 both produce an
irfftn input whose last axis is n // 2 + 1 == 4, so the requested size is
unrecoverable from the primitive state by the time vmap runs. Fixing it means
adding an output size field, which changes state() and therefore export
serialization, so I left it out of this change.

Validation

M3 Pro, macOS 26.x, against main @ e78d894. Identical results on GPU and with
mx.set_default_device(mx.cpu), as expected for a graph level change.

check	result
every axis subset and order, 1-D to 4-D, all eight fft variants, linearity identity jvp(f)(x,t) == f(t) and the vjp transpose identity Re<L t, c> == Re<t, L* c>	1128 checks, 0 failures per device
vmap(jvp), jvp(vmap), jvp(jvp) and vjp(jvp) (the last two exactly zero, as required for a linear map)	0 failures
norm= backward, ortho and forward under jvp and vmap	0 failures
fft over a size-1 axis, under jvp and vmap	0 failures
grad(vmap(f)) against per-slice grads, and vmap(grad(f))	0 failures
vmap against the stack of the op over each slice, five axis orders including reversed and positive	0 failures
mx.compile cold and warm, including compiled vmap	16 checks, 0 failures per device
export and import round trip	pass
200 consecutive runs of the new tests	0 failures
python suite, GPU and cpu	815 tests, pre-existing test_fft_too_large on cpu only
C++ suite, both devices	263 cases, 3577 and 3581 assertions

Not verified locally: CUDA, Linux, Windows.

Tests

test_fft_vmap checks vmap against the stack of the op applied to each slice,
and test_fft_jvp checks the tangent against f(t), both over real and inverse
and five axis orders. Both fail on main.

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)