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

[michaelellis003]

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)

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

• [Feature Request] Multinomial with probabilities sampling #2212 (multinomial sampling from probabilities) — the inverse-CDF
  path suggested here is one efficient way to serve that use case at
  scale.