Compile spectral descent Newton-Schulz with foreach_map

[mlazos]

Summary

• express the BF16 Newton-Schulz spectral descent update with foreach_map
• compile the foreach preconditioner separately when Shampoo PT2 compilation is enabled
• retain the existing implementation for other dtypes and SVD orthogonalization

The separate compiled callable is needed because the surrounding optimizer preconditioning path intentionally disables compilation. It lets compatible matrix updates reach PT2 as one operation and enables grouped symmetric GEMM lowering when used with the corresponding PyTorch Inductor changes.

This draft currently depends on PyTorch Inductor symmetric/grouped GEMM changes that are not yet upstream.

Test plan

ruff check distributed_shampoo/distributed_shampoo.py distributed_shampoo/preconditioner/spectral_descent_preconditioner_list.py distributed_shampoo/preconditioner/tests/spectral_descent_preconditioner_list_test.py
ruff format --check distributed_shampoo/distributed_shampoo.py distributed_shampoo/preconditioner/spectral_descent_preconditioner_list.py distributed_shampoo/preconditioner/tests/spectral_descent_preconditioner_list_test.py
python -m pytest distributed_shampoo/preconditioner/tests/spectral_descent_preconditioner_list_test.py -q -k 'bfloat16_uses_foreach_map or empty_list'

The two focused tests pass. Running the entire test file gives 11 passes and two unrelated SVD failures because the locally built PyTorch does not include LAPACK.

Local performance measurements with the dependent PyTorch changes:

• four 1536 x 5120 BF16 matrices: 2.043 ms to 1.396 ms for the isolated preconditioner (1.46x)
• two-rank Gloo full optimizer step on two GB200 GPUs: 31.386 ms baseline and 31.992 ms compiled (0.981x)

The distributed result is included to avoid implying an end-to-end speedup; communication dominates this configuration. An NCCL-enabled measurement is still needed.