Add inverse-CDF path for large categorical sampling

[MicroMilo]

PR: 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, complex logits, other overloads, and smaller workloads remain
  unchanged.
• Keep the source helpers in random_impl.h rather than adding experimental
  functions to the public C++ random API.

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	23.70x	[23.17, 24.56]	75.12x
B=1, N=M=8000	156.56x	[148.44, 163.51]	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

• Metal Python suite: 777 passed, 3 skipped, 10,877 subtests passed.
• CPU-only categorical/export subset: 35 passed, 1 Metal-only skipped, 6
  subtests passed.
• 4,096 randomized CPU/Metal/oracle cases, including 2,048 strided cases: zero
  failures.
• Debug -Werror C++ suite excluding unrelated linalg failures: 246/246 passed.
• ASan C++ suite excluding linalg: 230/230 passed; dedicated uint64 boundary
  test also passed under UBSan with halt_on_error=1.

[angeloskath]

I think there are more general ways to go about it. This is simply a searchsorted in the CDF so I would implement the searchsorted op first and then implement the inverse CDF sampling for the cases where we have a large number of samples.

I would advise to also split the PRs as they will be easier to review and merge.