Compute exact matrix powers in matrix_inverse_pth_root for any p

[TheSaiEaranti]

Problem

optax.matrix_inverse_pth_root's docstring promises matrix^(-1/p) for any positive integer p, but the internal mat_power helper selects an unrolled matrix power via round(log2(p)), which is exact only for p ∈ {1, 2, 4, 8}. For every other p the coupled Newton iteration converges to the wrong fixed point while still reporting a tiny error indicator:

 p   max rel err vs eigh M^(-1/p)   internal err
 3                    1.857e-01       4.768e-07   ← silent
 5                    1.073e-01       2.980e-07   ← silent
 6                    7.434e-02       5.960e-07   ← silent
 7                    3.017e-02       2.384e-07   ← silent
 (1, 2, 4, 8 are all ~1e-7)

The failure is independent of any reference implementation: for p=3, inv(result³) differs from the input matrix by 12% relative error. p=6 arises in practice as 2 * ndim when preconditioning rank-3 tensors Shampoo-style. Nothing surfaces the restriction — no error, no docstring note — and the existing test only exercises p=4.

Fix

Replace the rounded-power switch with exponentiation by squaring in a lax.while_loop — exact for any positive integer p, including traced values. This mirrors what the upstream google-research/scalable_shampoo implementation did when it replaced this same construction with an exact mat_power. All p in 1..8 now match the eigendecomposition reference to ~1e-7, on and off jit.

One trade-off worth flagging: the loop does ~2 matmuls per bit of p (e.g. 6 for p=8 versus 3 unrolled before). If that matters for Shampoo-scale workloads, a concrete-int fast path that unrolls exactly at trace time is easy to layer on — happy to add it here if preferred.

Tests

test_matrix_inverse_pth_root_arbitrary_p, parameterized over p ∈ 1..8, asserts both the returned error indicator and the max relative error against the eigh-based M^(-1/p). On the previous implementation exactly the four silent values (3, 5, 6, 7) fail and the four correct ones pass. Full linear_algebra_test.py: 31 passed, 240 subtests.