[preconditioner] thread-safe TF32 context manager, generalize rank_deficient_stability_config to MatrixFunctionConfig, refactor _precondition_grad

[Vishal-sys-code]

Resolves three open TODOs in distributed_shampoo/preconditioner.

1. Thread-safe TF32 context manager: _matrix_inverse_root_coupled_higher_order_newton was toggling torch.backends.cuda.matmul.allow_tf32 directly with try/finally. Since this flag is process-global, concurrent calls from multi-threaded optimizer setups could race on it. Added _TF32_LOCK (RLock) and _scoped_tf32_setting context manager so the read/modify/restore is atomic. The fast path (when disable_tf32 is False) skips lock acquisition entirely. Added a multithreaded regression test in matrix_functions_test.py.

2. Generalize rank_deficient_stability_config to MatrixFunctionConfig: The field was on EigendecompositionConfig with a TODO to move it up to the base class. Moved it to MatrixFunctionConfig along with the default factory. Cleaned up the conditional TODOs in PerturbationConfig and PseudoInverseConfig docstrings that were contingent on this generalization. Removed the assert isinstance(..., EigendecompositionConfig) guards in _compute_preconditioned_gradient and _compute_outer_product_list that only existed because the field was not on the base class yet.

3. Refactor _precondition_grad: The implementation used functools.reduce with a lambda that called grad.to(dtype=target_dtype) on every iteration and reconstructed (*range(1, grad.ndim), 0) on every non-preconditioned step. Replaced with a direct loop, a single upfront dtype cast, and a precomputed permutation tuple. Removed the now-unused functools.reduce import.