v0.32.1

What's Changed

• Fix int64 type cast error when loading GGUF metadata arrays by @danlee2002 in #3823
• Document default values in normalization layer docstrings by @Pablosinyores in #3819
• Warn at configure time when NAX kernels are disabled by @pierre427 in #3824
• [CUDA][Improvement] RMSNorm forward speed up by @nastya236 in #3850
• Fix captured random state in compile by @angeloskath in #3828
• Fix JIT preamble header filter matching project paths containing "Xcode" by @apocryphx in #3873
• Document default value of p in dropout layer docstrings by @ayaangazali in #3870
• [WIP] [CUDA] fsdp by @nastya236 in #3768
• Fix triplet_loss docstring to document the reduced output shape by @ayaangazali in #3884
• Zero-copy CPU import: mx.array(host_buffer, copy=False) on unified memory by @HaoXuAI in #3872
• Round MLX_SDPA_BLOCKS up to a multiple of 32 by @pierre427 in #3875
• Reuse Metal WAR tracking hash tables by @neilmehta24 in #3882
• Use unroll_count(4) for the NAX attention Q@K.T loop by @wyanzhao in #3843
• metal: add gemv_wide for fp16/bf16 matmuls of a few rows by @jessegross in #3888
• Fix broken docstring rendering in Linear and RNN by @ayaangazali in #3890
• Fix Adamax betas docstring and MultiOptimizer filters type by @ayaangazali in #3889
• Fix incorrect nvfp4 quantized_matmul through the split-K path by @metascroy in #3854
• [Metal] Avoid regex in custom kernel name generation by @aleroot in #3869
• Fix prod dtype promotion when reducing a size-1 axis by @eyupcanakman in #3898
• [CUDA] Fix grid overflow in gemm conv unfold kernels for >= 65,536 output positions by @AdamDLuz in #3893
• [CUDA] columnwise quantize with tma by @nastya236 in #3157
• metal: reduce NVFP4 scales per 16-lane group by @jessegross in #3934
• Update homebrew in CI by @angeloskath in #3946
• Make index autodiff errors explicitly recommend stop_gradient by @dogukanveziroglu in #3820
• Making JACCL coordinator optional by @angeloskath in #3899
• Fix docstring mismatches in the Python bindings by @ayaangazali in #3948
• Raise a clear error for an invalid quantization mode in nn layers by @ayaangazali in #3914
• Fix incorrect examples and outputs in the usage docs by @ayaangazali in #3956
• Skip test_gather_qmm_sorted cpu test on M1 mac by @zcbenz in #3973
• Fixes an axis mismatch bug in matrix norm for case -1 and 1 by @danlee2002 in #3827
• Fix BatchNorm running variance estimator by @ishtihoss in #3817
• Fix build error caused by TMA macro guard by @zcbenz in #3988
• Fix Glorot/He uniform init docstrings to label the uniform bound, not sigma by @vineethsaivs in #3831
• Fix custom metal kernel cache collision for same name, different source by @katlun-lgtm in #3833
• Fix log_cosh_loss docstring to document the element-wise loss by @winklemad in #3846
• Fix step activation docstring to match >= threshold behavior by @ayaangazali in #3902
• docs: document the reduced-precision float32 default and MLX_ENABLE_TF32 by @stoyoda0012-cyber in #3894
• Fix InstanceNorm Shape docstring to require at least 3 dimensions by @ayaangazali in #3903
• Fix filter_and_map docstring argument order for filter_fn and is_leaf_fn by @ayaangazali in #3906
• Export C++20 requirement to CMake consumers by @PhysicistJohn in #3971
• Fix implicit thread address space qualifier becoming explicit in metal 4.1 by @louen in #3963
• Fix Transformer ignoring a custom encoder or decoder with no parameters by @ayaangazali in #3962
• docs: remove references to removed --no-verify-script launch flag by @latent-9 in #3959
• added eye(0) support and tests by @aaishwarymishra in #3952
• Refactor the JACCL ring and add threads for the multiple rings by @angeloskath in #3900
• Fix shapeless matmul with dynamic batch dimensions by @varshneydevansh in #3813
• Fix JIT build with old macOS SDK by @metascroy in #3853
• Fix cholesky_inv documented argument name to match the binding by @ayaangazali in #3950
• Add new_thread_unsafe_stream to the devices and streams docs by @ayaangazali in #3968
• Add gather_qqmm by @zcbenz in #3757
• Fix missing printoptions doc page from autosummary filename collision by @ayaangazali in #3985
• Document ThreadLocalStream and iinfo in the API reference by @ayaangazali in #3986
• Make "stop" optional in arange by @aaishwarymishra in #3982
• Add softsign to the nn functions docs by @ayaangazali in #3989
• Fix broken all_sum and Group references in the data parallelism example by @ayaangazali in #3996
• Fix ast.metal_kernel typo in the custom Metal kernels guide by @ayaangazali in #3997
• Fix unresolved mx.array docstring references by @ayaangazali in #3990
• python: fix bfloat16 buffer format itemsize mismatch by @reckylurker in #3975
• Template Metal complex scalar lanes by @PhysicistJohn in #3970
• Pad 2D conv input channels to reach the specialized Metal kernel by @eyupcanakman in #3904
• Support head dimension 96 in Metal full attention by @dhiltgen in #3943
• Fix mx.remainder floored-mod for float16/bfloat16 on CPU by @sashko-zakharchuk in #3976
• Fix sorted gather_mm activation row stride by @metascroy in #3960
• Fix state corruption when a primitive throws during eval by @WindChimeRan in #3675
• added complex support by @aaishwarymishra in #3984
• Validate freeze and unfreeze keys against the whole model when recursing by @ayaangazali in #3966
• Fix mlx.launch --python: flag is parsed but never forwarded to the launch script by @jonathan308 in #4002
• Fix broken fully_shard reference in FullyShardedModule docstring by @ayaangazali in #4007
• Use the current interpreter in the comparative benchmark runner by @ayaangazali in #4016
• Document MLX environment variables by @XXXXRT666 in #4000
• Fix CUDA batched GEMV grid overflow by @jasp-nerd in #3929
• Fix all_gather benchmark collapsing its input to a scalar by @ayaangazali in #4017
• Check threadgroup size in the 1-pass sdpa_vector dispatch by @apocryphx in #4018
• Align example projects with the Python 3.10 minimum by @ayaangazali in #4024
• Fix C++ benchmarks failing to build on overloaded astype by @ayaangazali in #4025
• Fix conv_transpose maxBufferLength failures on Metal via tiled unfold by @eyupcanakman in #3845
• Give each host a unique rank in Hostfile.from_list by @ayaangazali in #4027
• Fix crash when reporting partial rings in mlx.distributed_config by @ayaangazali in #4026
• docs: Do not pass MLX_METAL_FAST_SYNCH=1 by default by @katlun-lgtm in #4005
• Fix signed-integer overflow in convolution shape arithmetic by @eyupcanakman in #3938
• Template Metal C2C FFT scalar lanes by @PhysicistJohn in #3969
• Fix segfault in expand_dims for out of bounds negative axes by @Gusanidas in #4021
• Add optional dtype parameter to zeros_like and ones_like by @reckylurker in #4028
• Fix mx.longsumexp output shape issue. by @JasonHonKL in #4030
• Treat backend and envs as optional when parsing a hostfile by @ayaangazali in #4039
• Report bad hostfiles as CLI errors in mlx.launch by @ayaangazali in #4040
• Fix ceil error message to say ceil instead of floor by @PhysicistJohn in #4042
• Fix installed static MLX package on Windows by @yingjiacai in #3848
• Derive the qmv fast path K alignment from bits by @jlee5814 in #3965
• Fix eigh UPLO and zero-size eigh/svd on the CPU by @kapellirohith in #3834
• Add metadata to exported functions by @eyupcanakman in #3855
• Fix empty matrix case in cholesky by @prady0t in #4033
• Reject a negative max_norm in clip_grad_norm by @ayaangazali in #4058
• Build extensions with the interpreter running the build by @ayaangazali in #4057
• Reject negative dimensions in broadcast_to and random shapes by @ayaangazali in #4046
• Report bad host arguments as CLI errors instead of tracebacks by @ayaangazali in #4045
• Fix Metal FFT for sizes above 2**20 by @adityasingh2400 in #4013
• Request MSL 4.1 from the runtime compiler on macOS 27 by @magnuslundstedt in #4052
• Fix concurrent Metal kernel cache lookup by @kitty-eu-org in #4043
• Refactor CPU ArgReduce dtype dispatch by @PhysicistJohn in #4066
• Refactor NumberOfElements dtype dispatch by @PhysicistJohn in #4064
• Refactor CPU Select dtype dispatch by @PhysicistJohn in #4065
• Raise qmv batch limit for large matrices on M5-class GPUs by @pierre427 in #3791
• Enable half-precision complex Metal FFT kernels by @PhysicistJohn in #3981
• Fix: mx.distributed.sum_scatter crashes on a scalar(0-d) by @JasonHonKL in #4071
• Return an empty result from pinv for zero-size inputs by @ayaangazali in #4069
• Handle zero-size inputs in eig and qr on the CPU by @ayaangazali in #4068
• Normalize biases before encoding in gather_qmm_rhs by @dudududukim in #4056
• Keep randint samples inside [low, high) by @adityasingh2400 in https://github.com//pull/4012
• Allow saving empty arrays to npy and safetensors by @ayaangazali in #4080
• Fix DeviceType annotations for device and stream arguments by @XXXXRT666 in #4059
• Refactor CPU DivMod dtype dispatch by @PhysicistJohn in #4088
• Refactor CPU SegmentedMM dtype dispatch by @PhysicistJohn in #4087
• Refactor CPU Arange dtype dispatch by @PhysicistJohn in #4085
• Fix Device and Stream lexicographic ordering by @wgu9 in #4086
• Enable complex64 scatter addition on GPU by @PhysicistJohn in #4078
• Optimize large NVFP4 QMV on M5 Max by @dhiltgen in #3961
• ring: fail on peer disconnect instead of hanging forever by @erwinzhang7 in #4060
• Fix: propagate NaN when dealing with cummax and curmin by @JasonHonKL in #4047
• Refactor CPU equal_nan dtype dispatch by @PhysicistJohn in #4095
• Refactor CPU Matmul dtype dispatch by @PhysicistJohn in #4103
• Refactor CPU Gather dtype dispatch by @PhysicistJohn in #4105
• Refactor generic CPU unary dtype dispatch by @PhysicistJohn in #4106
• Refactor generic CPU binary dtype dispatch by @PhysicistJohn in #4107
• Validate pooling kernel size, stride and padding by @ayaangazali in #4101
• Only reject max and min when a reduced axis is empty by @ayaangazali in #4079
• Fix SinusoidalPositionalEncoding silently overriding explicit scale=0.0 by @YassineMA03 in #4098
• Fix any and all treating -0.0 as nonzero by @ayaangazali in #4090
• Fix Eigh CPU dtype error message by @PhysicistJohn in #4130
• Fix random distribution dtype error messages by @PhysicistJohn in #4131
• Fix safetensors dtype error formatting by @PhysicistJohn in #4132
• Fix no-GPU synchronize error formatting by @PhysicistJohn in #4133
• Fix: mx.isinf doens't consider imaginary plane. by @JasonHonKL in #4092
• Refactor CPU Copy destination dtype dispatch by @PhysicistJohn in #4104
• Refactor CPU GatherAxis dtype dispatch by @PhysicistJohn in #4109
• Refactor inexact CPU unary dtype dispatch by @PhysicistJohn in #4111
• Refactor CPU comparison dtype dispatch by @PhysicistJohn in #4112
• Replace Ring's local dtype macro with common dispatch by @PhysicistJohn in #4120
• Refactor CPU ArgSort dtype dispatch by @PhysicistJohn in #4123
• Refactor CPU Softmax dtype dispatch by @PhysicistJohn in #4127
• Use common dtype dispatch for Metal Arange by @PhysicistJohn in #4128
• Pick BM from rows per expert in gather_qmm_rhs_nax by @dwijenpatel in #4023
• Enable complex64 take_along_axis backward on Metal by @PhysicistJohn in #4094
• Add mx.searchsorted with CPU, Metal and CUDA kernels by @erwinzhang7 in #4035
• Stop integer power from hanging on a negative exponent by @ayaangazali in #4100
• Fix second order gradients for sort, partition, topk and cummax/cummin by @Adityaj0 in #4117
• Fix segfault on negative out of bounds axes in take_along_axis/put_along_axis by @Adityaj0 in #4118
• Deduplicate slow CPU convolution dtype dispatch by @PhysicistJohn in #4126
• Use common dispatch for compiled constants by @PhysicistJohn in #4129
• Fix set_printoptions example by @PhysicistJohn in #4140
• Clarify export output type diagnostic by @PhysicistJohn in #4141
• Fix integer keys collision in tree_unflatten by @Solaris-star in #3878
• Fix fence tracking for donated dynamic slice offsets by @michaelellis003 in #4099
• Deduplicate distributed all-reduce construction by @PhysicistJohn in #4144
• Deduplicate distributed socket errors by @PhysicistJohn in #4145
• Remove redundant Metal bitwise dispatch by @PhysicistJohn in #4147
• Deduplicate custom Metal kernel size validation by @PhysicistJohn in #4149
• Deduplicate CUBLAS compute type cases by @PhysicistJohn in #4151
• Refactor CPU Scatter dtype dispatch by @PhysicistJohn in #4155
• fix(cpu): name the right routine in LUF and SVD error messages by @devteamaegis in #4167
• Refactor CPU ArgPartition dtype dispatch by @PhysicistJohn in #4156
• add complex to scalar to avoid raising warnings by linters by @aaishwarymishra in #4168
• Refactor CPU Copy source dtype dispatch by @PhysicistJohn in #4170
• Deduplicate Metal scan encoder setup by @PhysicistJohn in #4154
• [Metal] Skip empty NAX GEMM output groups by @XXXXRT666 in #3941
• docs: fix stale CLI invocations in LLM inference example by @deBrian07 in #4172
• Increase ccache size in CI by @zcbenz in #3999
• Fix hadamard_transform on the GPU for n = m with no power-of-2 factor by @axiom-of-choice in #4054
• Refactor CPU Partition dtype dispatch by @PhysicistJohn in #4175
• Refactor CPU ScatterAxis and MaskedScatter dtype dispatch by @PhysicistJohn in #4176
• docs: document softmax's precise argument by @apocryphx in #4178
• Send out of range trig arguments to libm by @ayaangazali in #4157
• Add reflect and symmetric padding modes to mx.pad by @katlun-lgtm in #3608
• docs: add a usage page on growing arrays by @apocryphx in #4019
• Fix and enable non-transposed NAX qmm by @gordofreemo in #4051
• Add scatter reduce for JACCL by @angeloskath in #3901
• Fix diag for zero-size input by @devteamaegis in #4165
• Add an inverse-CDF path to categorical sampling by @deescondrillas in #4177
• Fail when a requested RDMA device is not found by @erwinzhang7 in #4180
• Bound GGUF tensor data offsets against the file mapping by @robertomeroni in #4179
• Shift by the max in log_softmax by @ayaangazali in #4169
• Shift by the max in cross_entropy by @ayaangazali in #4188
• Fix linalg.norm(x, keepdims=True) not keeping dims by @devteamaegis in #4166
• Refactor inexact CPU binary dtype dispatch by @PhysicistJohn in #4194
• Implement batched matmul for large 1D dot products by @Ved235 in #3580
• Re-disable qmm_n_nax and fix group_size < 64 by @PhilipJohnBasile in #4202
• fix: shell-quote the working directory in the distributed launch script by @bunlongheng in #4189
• Report when no usable GID is found by @erwinzhang7 in #4191
• Keep double precision for python floats in float64 operations by @nileshpatil6 in #4173
• Jaccl ring multi wire fix by @angeloskath in #4193
• Report why creating a queue pair failed by @erwinzhang7 in #4209
• Add antialias support to nn.Upsample for linear and cubic modes by @lyonsno in #3677
• Require pytorch 2.12 for metal dlpack tests by @zcbenz in #4206
• Remove std::optional in C++ public interface by @zcbenz in #4207
• Fix SliceUpdate JVPs with one traced input by @PhysicistJohn in #4200
• Raise instead of hanging when reflect/symmetric pad gets an empty axis by @codeAnqiang-ma in #4223
• Reuse Stockham twiddles in large batched Bluestein FFTs by @PhysicistJohn in #4084
• Validate decay_steps and step_size in the schedulers by @ayaangazali in #4217
• Add runnable distributed examples by @erwinzhang7 in #4219
• Break monolithic MTLResidencySet into smaller sets by @JamesMcCarthy44 in #4211
• [CUDA][Improvement] RMSNorm backward by @nastya236 in #3881
• Fix MultiOptimizer on models containing empty modules by @erenmenges in #4215
• Return tuple in meshgrid by @JasonHonKL in #4229
• Add endpoint parameter to linspace by @AKnassa in #4184
• Fix vmap of partition/argpartition dropping the kth argument by @Adityaj0 in #4116
• Fix nan_to_num replacing inf with 0 for float16 and bfloat16 by @codeAnqiang-ma in #4222
• Fix einsum not broadcasting batch dimensions in batched tensordot by @Adityaj0 in #4125
• Dequantize in float32 by @RohanGautam in #4241
• Reject complex in erf and erfinv by @ayaangazali in #4243
• Make abs a no-op for unsigned types in the simd helpers by @ayaangazali in #4240
• Fix quantize matrix multiplication floor issue by @JasonHonKL in #4251
• Only use MPI backend for world size > 1 by @erwinzhang7 in #4210
• Reject complex in expm1, sigmoid and arctan2 by @ayaangazali in #4257
• Decompose small kernel-depth 3D convs into 2D convs by @katlun-lgtm in #3785
• Fix Metal sort of a view with a negative stride by @robertomeroni in #4252
• Mirror the depth axis in the decomposed 3D conv when flipped by @ayaangazali in #4277
• Fix Metal row reductions on negative-stride views by @FU-max-boop in #4267
• [CUDA] Fix custom kernel cache collision for same name, different source by @HaoXuAI in #4273
• Fix ops rejecting integers larger than INT32_MAX by @FeliGame in #4255
• Fix var/std for complex numbers by @ayaangazali in #4260
• Fix int32 overflow in conv padded input and pad shapes by @eyupcanakman in #4258
• Reject complex in remainder by @ayaangazali in #4270
• Compare the macOS SDK version as a version when gating JACCL by @erwinzhang7 in #4286
• Clamp ring socket transfers so a payload of 2 GiB or more can be sent by @erwinzhang7 in #4281
• Replace hand-rolled axis normalization with normalize_axis_index in split/unstack/partition/topk by @Adityaj0 in #4288
• Remove grouped output in CI by @zcbenz in #4195
• [CUDA] Fix finding cuda 13 headers in JIT compilation by @zcbenz in #3995
• Refactor wheel building script by @zcbenz in #3818
• Make mx.compile cache erasing thread safe by @zcbenz in #4248
• Add builds for free-threaded python by @zcbenz in #3812
• Fix int32 overflow in concatenate/repeat/kron by @ayaangazali in #4303
• Widen list elements that do not fit in int32 to int64 by @ayaangazali in #4305
• Propagate CPU errors to events by @zcbenz in #3742
• Fix mx.arange dtype inference overflow regression by @shoutoutuoadi325 in #4324
• Add workflow to update pull request limit bypass list by @zcbenz in #4320
• Support head dimension 72 in Metal full attention by @dhiltgen in #4330

New Contributors

• @pierre427 made their first contribution in #3824
• @apocryphx made their first contribution in #3873
• @ayaangazali made their first contribution in #3870
• @HaoXuAI made their first contribution in #3872
• @neilmehta24 made their first contribution in #3882
• @wyanzhao made their first contribution in #3843
• @metascroy made their first contribution in #3854
• @aleroot made their first contribution in #3869
• @AdamDLuz made their first contribution in #3893
• @vineethsaivs made their first contribution in #3831
• @winklemad made their first contribution in #3846
• @stoyoda0012-cyber made their first contribution in #3894
• @PhysicistJohn made their first contribution in #3971
• @latent-9 made their first contribution in #3959
• @aaishwarymishra made their first contribution in #3952
• @reckylurker made their first contribution in #3975
• @sashko-zakharchuk made their first contribution in #3976
• @WindChimeRan made their first contribution in #3675
• @jonathan308 made their first contribution in #4002
• @jasp-nerd made their first contribution in #3929
• @Gusanidas made their first contribution in #4021
• @JasonHonKL made their first contribution in #4030
• @yingjiacai made their first contribution in #3848
• @jlee5814 made their first contribution in #3965
• @kapellirohith made their first contribution in #3834
• @prady0t made their first contribution in #4033
• @magnuslundstedt made their first contribution in #4052
• @kitty-eu-org made their first contribution in #4043
• @dudududukim made their first contribution in #4056
• @erwinzhang7 made their first contribution in #4060
• @YassineMA03 made their first contribution in #4098
• @dwijenpatel made their first contribution in #4023
• @Adityaj0 made their first contribution in #4117
• @Solaris-star made their first contribution in #3878
• @michaelellis003 made their first contribution in #4099
• @devteamaegis made their first contribution in #4167
• @deBrian07 made their first contribution in #4172
• @axiom-of-choice made their first contribution in #4054
• @gordofreemo made their first contribution in #4051
• @deescondrillas made their first contribution in #4177
• @robertomeroni made their first contribution in #4179
• @Ved235 made their first contribution in #3580
• @PhilipJohnBasile made their first contribution in #4202
• @bunlongheng made their first contribution in #4189
• @nileshpatil6 made their first contribution in #4173
• @codeAnqiang-ma made their first contribution in #4223
• @JamesMcCarthy44 made their first contribution in #4211
• @erenmenges made their first contribution in #4215
• @AKnassa made their first contribution in #4184
• @RohanGautam made their first contribution in #4241
• @FU-max-boop made their first contribution in #4267
• @FeliGame made their first contribution in #4255
• @shoutoutuoadi325 made their first contribution in #4324

Full Changelog: v0.32.0...v0.32.1