Stabilize reduced-precision InstanceNorm

[ternaus]

Proposed changes

Fixes #4228.

This PR stabilizes reduced-precision InstanceNorm without changing its public (N, ..., C) layout or output dtype:

• reshape each input to (N, C, spatial_voxels) and use the stable fused mx.fast.layer_norm reduction;
• restore the original channels-last shape before the optional affine transform;
• add an FP16 regression case whose direct variance overflows even though every input value is finite.

Why

The current FP16 path computes mx.var in FP16. A channel alternating between -512 and 512 therefore gets var=inf, and InstanceNorm collapses the channel to zero instead of returning values near -1 and 1.

mx.fast.layer_norm already performs the same reduction needed by InstanceNorm after the spatial dimensions are flattened. It avoids a full FP32 input cast and uses the optimized normalization kernel.

Performance

Apple M4, MLX 0.32.0 core, FP16 input, five warmups and 20 measured calls:

NDHWC shape	Existing p50	FP32-cast p50	This PR p50
1 x 32 x 64 x 64 x 30	0.785 ms	0.789 ms	0.196 ms
1 x 128 x 96 x 96 x 30	2.357 ms	5.145 ms	1.098 ms

The larger shape is the first-stage feature map from the 3D segmentation route that exposed the overflow.

Validation

• PYTHONPATH=python:python/tests python -m unittest test_nn.TestLayers
• PYTHONPATH=python:python/tests python -m unittest test_nn
• pre-commit run --all-files

Checklist

• 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; no public API or documented behavior changed)