Use QR in ProximalProjection

[YigitElma]

          I want to believe there is a way to do this but may need some additional changes to some tests. Benchmarks for this,

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     +0.15 +/- 9.10     | +8.40e-04 +/- 4.95e-02 |  5.45e-01 +/- 3.3e-02  |  5.45e-01 +/- 3.7e-02  |
 test_equilibrium_init_medres            |     -2.49 +/- 3.23     | -1.16e-01 +/- 1.51e-01 |  4.54e+00 +/- 5.8e-02  |  4.66e+00 +/- 1.4e-01  |
 test_equilibrium_init_highres           |     -0.12 +/- 1.95     | -7.13e-03 +/- 1.15e-01 |  5.87e+00 +/- 8.5e-02  |  5.87e+00 +/- 7.7e-02  |
 test_objective_compile_dshape_current   |     -0.88 +/- 3.60     | -3.70e-02 +/- 1.51e-01 |  4.17e+00 +/- 1.2e-01  |  4.21e+00 +/- 9.1e-02  |
 test_objective_compute_dshape_current   |     +3.03 +/- 4.06     | +1.57e-04 +/- 2.10e-04 |  5.33e-03 +/- 2.0e-04  |  5.17e-03 +/- 6.3e-05  |
 test_objective_jac_dshape_current       |     +1.88 +/- 7.02     | +8.16e-04 +/- 3.05e-03 |  4.43e-02 +/- 1.7e-03  |  4.35e-02 +/- 2.5e-03  |
 test_perturb_2                          |     -0.86 +/- 1.87     | -1.74e-01 +/- 3.79e-01 |  2.01e+01 +/- 3.0e-01  |  2.02e+01 +/- 2.3e-01  |
+test_proximal_jac_atf_with_eq_update    |    -11.98 +/- 1.23     | -5.31e+00 +/- 5.46e-01 |  3.90e+01 +/- 3.9e-01  |  4.44e+01 +/- 3.8e-01  |
 test_proximal_freeb_jac                 |     -1.28 +/- 1.26     | -9.58e-02 +/- 9.47e-02 |  7.40e+00 +/- 4.2e-02  |  7.50e+00 +/- 8.5e-02  |
 test_solve_fixed_iter_compiled          |     +0.53 +/- 2.79     | +1.10e-01 +/- 5.85e-01 |  2.11e+01 +/- 5.1e-01  |  2.10e+01 +/- 2.9e-01  |
 test_LinearConstraintProjection_build   |     +1.59 +/- 2.14     | +1.73e-01 +/- 2.32e-01 |  1.11e+01 +/- 2.3e-01  |  1.09e+01 +/- 3.6e-02  |
 test_build_transform_fft_midres         |     +0.31 +/- 2.92     | +1.84e-03 +/- 1.75e-02 |  6.02e-01 +/- 1.5e-02  |  6.00e-01 +/- 9.8e-03  |
 test_build_transform_fft_highres        |     -0.33 +/- 2.30     | -3.20e-03 +/- 2.20e-02 |  9.52e-01 +/- 2.0e-02  |  9.56e-01 +/- 8.1e-03  |
 test_equilibrium_init_lowres            |     +0.12 +/- 0.92     | +4.87e-03 +/- 3.73e-02 |  4.07e+00 +/- 3.2e-02  |  4.06e+00 +/- 2.0e-02  |
 test_objective_compile_atf              |     -0.16 +/- 1.34     | -1.30e-02 +/- 1.11e-01 |  8.27e+00 +/- 8.4e-02  |  8.28e+00 +/- 7.3e-02  |
 test_objective_compute_atf              |     -0.19 +/- 1.37     | -3.01e-05 +/- 2.15e-04 |  1.57e-02 +/- 1.5e-04  |  1.58e-02 +/- 1.5e-04  |
 test_objective_jac_atf                  |     +1.68 +/- 2.13     | +3.27e-02 +/- 4.14e-02 |  1.98e+00 +/- 3.2e-02  |  1.95e+00 +/- 2.7e-02  |
 test_perturb_1                          |     +0.10 +/- 1.02     | +1.41e-02 +/- 1.48e-01 |  1.45e+01 +/- 1.3e-01  |  1.45e+01 +/- 7.6e-02  |
+test_proximal_jac_atf                   |    -32.47 +/- 0.61     | -2.69e+00 +/- 5.09e-02 |  5.59e+00 +/- 3.1e-02  |  8.28e+00 +/- 4.1e-02  |
+test_proximal_freeb_compute             |    -46.95 +/- 0.56     | -9.54e-02 +/- 1.13e-03 |  1.08e-01 +/- 6.5e-04  |  2.03e-01 +/- 9.2e-04  |
 test_solve_fixed_iter                   |     +0.34 +/- 1.77     | +1.07e-01 +/- 5.65e-01 |  3.20e+01 +/- 2.9e-01  |  3.19e+01 +/- 4.9e-01  |

Originally posted by @YigitElma in #1409 (comment)

In _proximal_jvp_f_pure, we use SVD, which is what makes up over 90% of the time spent. It should be possible to use similar QR implementation. I checked the matrix that we try to invert in test_proximal_jacobian and its condition number is actually pretty low, so I don't think we need much regularisation. The bad part is Gxh.

[YigitElma]

Even if we don't change it, I think it would help for future devs to put some comments why we use SVD over other methods. For example, in perturb, I think SVD is overall better, because it requires 1 decomposition and then trust_region_exact_svd is basically free but trust_region_exact_qr is very expensive in most cases. Also, we call trust region subproblem solver multiple times for the same matrix.

[f0uriest]

what did you use to check the condition number? In general the condition number varies a ton across the optimization space and depends on the particular combination of objectives. The regularization in the SVD was added to reduce stalling out like in #802