Add inverse-CDF path for large categorical sampling

[MicroMilo]

Add a transform-aware inverse-CDF path for large categorical sampling

Summary

This change adds a source-level CategoricalSearch primitive and uses it to avoid
the O(B*N*M) Gumbel temporary created by
random::categorical(logits, num_samples=M) for large Metal workloads.

Fixes #3847.

The inverse-CDF path uses FP32 softmax probabilities rounded onto a 52-bit uint64
mass grid, an exact uint64 cumulative sum, one uint32 random word per output
sample, and a Metal binary search kernel. CPU and non-Metal backends retain the
existing Gumbel-max path.

Why

Issue #3847 reports that drawing many samples from a large categorical creates a
Gumbel tensor with B*N*M elements. For N=M=1,000,000, that formulation would
require roughly 4 TB for the FP32 temporary.

A float32 inverse-CDF prototype reduced memory but failed an adversarial support
test: with one million categories and a 0.99 head probability, only 249,772 CDF
entries had positive increments. The uint64 fixed-point representation preserves
all one million positive weights.

Implementation

• Add CategoricalSearch(cdf[..., N], bits[..., M]) -> uint32[..., M].
• Add exact CPU evaluation for oracle testing.
• Add a Metal kernel using binary search and an overflow-safe 32x64 target map.
• Add an explicit vmap rule for both-mapped, CDF-only, and random-bits-only cases.
• Add export/import serialization.
• Normalize non-contiguous inputs inside backend evaluators so Contiguous does
  not leak into exported graphs.
• Define Metal-compatible non-finite behavior before the fixed-point cast:
  first +inf wins, NaNs are masked when another valid category exists, and
  rows containing only NaN/-inf choose category zero.
• Dispatch only on Metal for the num_samples overload when B*N*M >= 2^19.
  CPU, CUDA, other overloads, and smaller workloads remain unchanged.

Correctness

• CPU and Metal are bitwise identical for the same explicit CDF/random words.
• 1D/2D/3D batch and first/middle/last category axes pass.
• FP32/FP16/BF16/integer/bool logits follow an explicit FP32 softmax policy.
• direct, compile, vmap, compile-vmap, nested batched vmap, and export/import pass.
• Three-seed distribution test: max z-score 2.18, max TV 0.00139.
• One-million-category skew stress:
  • positive float CDF steps: 249,772;
  • positive fixed weights: 1,000,000;
  • fixed quantization TV: 4.94e-6;
  • tail z-score: 0.20.
• One-million-category uniform stress: max 32-bucket z-score 2.83.

Performance (Apple M5)

31 interleaved paired measurements, bootstrap 95% CI:

Workload	Median speedup	95% CI	Peak-memory reduction
B=8, N=M=1024	26.24x	[25.37, 26.89]	75.12x
B=1, N=M=8000	158.85x	[151.37, 172.52]	530.94x

The final 2^19 guard selected 23 formal grid cases; all 23 had a paired-speedup
CI lower bound above 1.02. Source graph checks confirm that guard-outside and CPU
paths contain the stock ArgReduce graph and do not construct
CategoricalSearch.

Compatibility and limitations

• The inverse-CDF path is deterministic for a given key but does not preserve
  Gumbel-max's exact random sequence.
• The performance threshold was calibrated on an Apple M5 and may need broader
  device validation.
• CUDA retains the stock implementation; no CUDA search kernel is included.
• The fixed-point representation preserves FP32 softmax support through the
  cumulative sum but does not provide higher-precision probabilities than FP32.

Tests

• python/tests/test_random.py: 17 passed.
• python/tests/test_export_import.py: 19 passed, 6 subtests passed.
• CPU-only, no-Metal/no-CUDA static MLX build completed successfully.