`mx.random.categorical(logits, num_samples=M)` allocates O(N·M) memory

## Summary `mx.random.categorical` with `num_samples=M` materializes an intermediate of shape `(M, N)` (one Gumbel perturbation per category per sample), so peak memory scales as O(N·M) rather than O(N + M). For the common case of drawing many samples from a large categorical — e.g. multinomial resampling in a particle filter, where N = M = number of particles — this makes the op unusable at scale: at N = M = 10⁶ it attempts a ~4 TB allocation. ## Reproduction (mlx 0.32.0, Apple M3 Pro) ```python import mlx.core as mx for n, m in [(2000, 2000), (4000, 4000)]: mx.reset_peak_memory() logits = mx.zeros((n,)) idx = mx.random.categorical(logits, num_samples=m, key=mx.random.key(0)) mx.eval(idx) print(f"N={n} M={m}: peak {mx.get_peak_memory() / 1e6:.1f} MB") ``` ``` N=2000 M=2000: peak 16.0 MB N=4000 M=4000: peak 64.1 MB ``` Peak memory quadruples when N and M each double — the O(N·M) signature (a 4000×4000 float32 intermediate is 64 MB). Extrapolated, N = M = 10⁶ needs 10¹² float32 ≈ 4 TB, which raises a `RuntimeError` on the allocation. ## Why it matters Drawing M samples from an N-way categorical is a core primitive for resampling (SMC / particle filters), bootstrap resampling, and any categorical-with-replacement sampling. The natural formulation `categorical(logits, num_samples=M)` is O(N·M) today, so users must avoid the built-in and hand-roll an inverse-CDF sampler (`cumsum(softmax(logits))` + a binary search over M sorted or systematic uniforms), which is O(N + M) in memory and works at N = M = 10⁶ in well under a gigabyte. ## Suggestion Either (a) document the O(N·M) memory characteristic so users know to avoid it for large N·M, or (b) offer an inverse-CDF path (or let `categorical` dispatch to one) when `num_samples` is large relative to the memory budget. Happy to share the inverse-CDF resampling kernel we use downstream if useful. ## Environment - mlx 0.32.0 - Apple M3 Pro, macOS 26.2 ## Related - #2212 (multinomial sampling from probabilities) — the inverse-CDF path suggested here is one efficient way to serve that use case at scale.