Fix segfault on negative out of bounds axes in take_along_axis/put_along_axis

[Adityaj0]

Fixes #4115.

Proposed changes

The bounds check in take_along_axis, scatter_axis (which backs put_along_axis and scatter_add_axis) and linalg::cross was dead code:

if (axis + a.ndim() < 0 || axis >= static_cast<int>(a.ndim()))

array::ndim() returns size_t, so axis + a.ndim() is evaluated in unsigned arithmetic and is never negative. Only the positive half of the check, which casts to int, actually ran.

A negative out of bounds axis therefore passed validation and reached broadcast_arrays as axis - int(a.ndim()). For small magnitudes that surfaced as an internal IndexError: SmallVector out of range; for larger ones it overwrote the stack and segfaulted inside broadcast_shapes:

Exception Type:  EXC_BAD_ACCESS (SIGSEGV)
Exception Subtype: KERN_INVALID_ADDRESS at 0x6d7ad7656d7ad74d

0  libmlx.dylib  mlx::core::broadcast_shapes(SmallVector<int, 10ul> const&, SmallVector<int, 10ul> const&) + 348
1  libmlx.dylib  mlx::core::BroadcastAxes::output_shape(...) + 576
2  libmlx.dylib  mlx::core::broadcast_arrays(...) + 196
3  libmlx.dylib  mlx::core::take_along_axis(...) + 348

This switches all three call sites to normalize_axis_index, which does the comparison in int and is already used for exactly this purpose elsewhere in ops.cpp (squeeze, expand_dims, flip, concatenate, stack, repeat, ...). Since it also normalizes, the later manual "allow negative axis" adjustments become redundant and are removed.

#4021 fixed this same pattern, but only in squeeze and expand_dims. It did
not touch these three call sites, which still carry the original expression. On
current main (which already includes #4021):

call	on main today
mx.expand_dims(x, -9)	ValueError: [expand_dims] Axis -9 is out of bounds ... (fixed by #4021)
mx.squeeze(x, -9)	ValueError: [squeeze] Axis -9 is out of bounds ... (fixed by #4021)
mx.take_along_axis(x, i, axis=-4)	IndexError: SmallVector out of range.
mx.put_along_axis(x, i, v, axis=-4)	IndexError: SmallVector out of range.
mx.linalg.cross(a, b, axis=-4)	IndexError: SmallVector out of range.

So this is a follow-up to #4021 covering the remaining call sites, not a
duplicate of it.

Before, on main:

take_along_axis(x, i, axis=3)     ValueError: [take_along_axis] Received invalid axis ...
take_along_axis(x, i, axis=-4)    IndexError: SmallVector out of range.
take_along_axis(x, i, axis=-100)  Segmentation fault
put_along_axis(x, i, v, axis=-4)  IndexError: SmallVector out of range.
linalg.cross(a, b, axis=-4)       IndexError: SmallVector out of range.

After, all of them raise the same clean error, e.g.

ValueError: [take_along_axis] Axis -100 is out of bounds for array with 3 dimensions.

Valid negative axes (-1, -2, -3) continue to work unchanged.

Note this slightly changes the wording of the existing error message for the positive out of bounds case, from "Received invalid axis for array with N dimensions" to normalize_axis_index's "Axis A is out of bounds for array with N dimensions", which also has the benefit of naming the offending axis. No tests depended on the old wording.

Tests

Added test_along_axis_invalid_axis and test_cross_invalid_axis to python/tests/test_ops.py, covering positive and negative out of bounds axes for take_along_axis, put_along_axis and linalg.cross, plus an assertion that valid negative axes still work. The deep-negative values are excluded from the test bodies only in the sense that they are asserted to raise; they crash on unpatched builds, which is the point.

• 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_ops, test_linalg, test_autograd and test_vmap pass. I also confirmed the new tests fail on an unpatched build.

[Adityaj0]

@zcbenz please review this PR. All checks passed.

[zcbenz]

For most open source projects reviewing contributor PR is a distraction that slows down actual development.