MLX: on-device token sampling with Gumbel-max#20454

Merged

metascroy merged 6 commits into

pytorch:mainfrom

kiymetakdemir:mlx-ondevice-sampling

Jun 25, 2026

Commits on Jun 23, 2026

MLX: on-device token sampling (mlx::sample, Gumbel-max)
kiymetakdemir
committed
MLX: reject non-uint32 width; move sample to custom_ops; docs
kiymetakdemir
committed

Commits on Jun 24, 2026

MLX: top-p sampling; temperature=0 greedy via IfNode; fp32 sampling chain; reorg sample tests
kiymetakdemir
committed

Commits on Jun 25, 2026