math: scan solve_pgs iterations

[staticpayload]

Summary

• replace the Python for-loop in solve_pgs with a lax.scan over iterations
• add a diagonal system test to validate projected solution

Testing

• python3.11 -m venv .venv-test
• pip install -e .
• pip install pytest
• pytest brax/math_test.py -q