Add benchmark for running fixed iteration solve

[f0uriest]

Should be merged before #1050 so benchmark will run.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 94.96%. Comparing base (0a8ec62) to head (9986152).
Report is 1879 commits behind head on master.

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #1054   +/-   ##
=======================================
  Coverage   94.96%   94.96%           
=======================================
  Files          87       87           
  Lines       21776    21776           
=======================================
+ Hits        20679    20680    +1     
+ Misses       1097     1096    -1     

see 1 file with indirect coverage changes

[github-actions]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     -1.76 +/- 3.64     | -9.99e-03 +/- 2.06e-02 |  5.57e-01 +/- 1.6e-02  |  5.67e-01 +/- 1.3e-02  |
 test_build_transform_fft_midres         |     -1.64 +/- 5.35     | -1.06e-02 +/- 3.46e-02 |  6.37e-01 +/- 2.9e-02  |  6.48e-01 +/- 1.9e-02  |
 test_build_transform_fft_highres        |     -1.50 +/- 1.99     | -1.57e-02 +/- 2.09e-02 |  1.03e+00 +/- 1.4e-02  |  1.05e+00 +/- 1.6e-02  |
 test_equilibrium_init_lowres            |     -1.84 +/- 2.06     | -7.72e-02 +/- 8.66e-02 |  4.12e+00 +/- 6.8e-02  |  4.19e+00 +/- 5.3e-02  |
 test_equilibrium_init_medres            |     -0.79 +/- 3.37     | -3.65e-02 +/- 1.56e-01 |  4.61e+00 +/- 1.4e-01  |  4.64e+00 +/- 6.4e-02  |
 test_equilibrium_init_highres           |     -0.63 +/- 1.93     | -3.82e-02 +/- 1.17e-01 |  6.03e+00 +/- 8.2e-02  |  6.07e+00 +/- 8.3e-02  |
 test_objective_compile_dshape_current   |     -0.77 +/- 1.75     | -3.08e-02 +/- 7.05e-02 |  3.99e+00 +/- 3.7e-02  |  4.02e+00 +/- 6.0e-02  |
 test_objective_compile_atf              |     -0.84 +/- 2.76     | -7.32e-02 +/- 2.39e-01 |  8.61e+00 +/- 1.9e-01  |  8.68e+00 +/- 1.5e-01  |
 test_objective_compute_dshape_current   |     +3.65 +/- 6.62     | +4.27e-05 +/- 7.75e-05 |  1.21e-03 +/- 6.9e-05  |  1.17e-03 +/- 3.6e-05  |
 test_objective_compute_atf              |     -2.15 +/- 8.50     | -9.32e-05 +/- 3.69e-04 |  4.24e-03 +/- 2.1e-04  |  4.33e-03 +/- 3.0e-04  |
 test_objective_jac_dshape_current       |     +2.77 +/- 6.71     | +1.04e-03 +/- 2.53e-03 |  3.87e-02 +/- 1.7e-03  |  3.76e-02 +/- 1.8e-03  |
 test_objective_jac_atf                  |     +0.91 +/- 4.45     | +1.69e-02 +/- 8.23e-02 |  1.86e+00 +/- 4.9e-02  |  1.85e+00 +/- 6.6e-02  |
 test_perturb_1                          |     -0.91 +/- 2.81     | -1.29e-01 +/- 3.98e-01 |  1.40e+01 +/- 1.6e-01  |  1.42e+01 +/- 3.7e-01  |
 test_perturb_2                          |     +3.02 +/- 2.13     | +5.82e-01 +/- 4.11e-01 |  1.99e+01 +/- 3.4e-01  |  1.93e+01 +/- 2.3e-01  |
 test_proximal_jac_atf                   |     +0.31 +/- 1.01     | +2.29e-02 +/- 7.46e-02 |  7.43e+00 +/- 4.7e-02  |  7.40e+00 +/- 5.8e-02  |
 test_proximal_freeb_compute             |     -0.56 +/- 1.30     | -9.96e-04 +/- 2.32e-03 |  1.78e-01 +/- 2.1e-03  |  1.79e-01 +/- 9.9e-04  |
 test_proximal_freeb_jac                 |     +0.87 +/- 1.44     | +6.51e-02 +/- 1.07e-01 |  7.51e+00 +/- 7.0e-02  |  7.44e+00 +/- 8.2e-02  |
 test_solve_fixed_iter                   |     +2.83 +/- 8.03     | +4.31e-01 +/- 1.22e+00 |  1.56e+01 +/- 9.5e-01  |  1.52e+01 +/- 7.7e-01  |

[YigitElma]

Actually, if we don't put jax.clear_caches() inside a setup function it will only be called once when this test runs. Do we want this or to clear cache before each solve? I assume we don't need to clear it every time since after first iteration of solve everything will be already jitted, but it is worth keeping in mind.

[f0uriest]

yeah we include clear_caches in the setup or run functions for some of the other tests (like the ones measuring compile time, or ones that are usually only run once like init) but in this case we're mostly interested in the average behavior over many iterations so it's not necessary.

I want to be sure this test tests correct thing ie clearing cache etc