Remove large arrays before big memory operations

[YigitElma]

This is a continuation of #1688.

• When we compute the new Jacobian, the old one is still in memory. This historically caused the second Jacobian computation to have a larger peak. This PR deletes the old Jacobian and matrix decomposition-related stuff before the new computation.
• Replaces the remaining non-in-place J *= d with a jitted scale_columns that donates its buffer, so the scaled matrix is written into J's memory; the old reference is deleted immediately. Some of these were noted in #1669 (comment).
• Operations like g = jnp.dot(J.T, f) create a copy of J. Since f is 1D, we can just compute g = jnp.dot(f.T, J) or equivalently g = jnp.dot(f, J) without explicitly forming J.T.
• The final J_h * 1 / d made two out-of-jit copies (J_h * 1, then / d). It now uses scale_columns too.

Some of these do not affect the peak memory, i.e. jnp.dot(J.T, f) change, but these are one-liners and can be useful in the future if they become the memory bottleneck. I also thought about updating select_step to take in J.T@J by using the decomposition stuff, but that is a bigger change and doesn't reduce peak memory.

[github-actions]

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |    0.39 %    |     4.219e+03      |     4.235e+03      |    16.26     |       27.55        |       23.59        |
  test_proximal_jac_w7x_with_eq_update   |   -0.41 %    |     6.805e+03      |     6.777e+03      |    -27.68    |       118.44       |       117.59       |
  test_proximal_freeb_jac                |    0.02 %    |     1.353e+04      |     1.354e+04      |     2.64     |       74.71        |       73.88        |
  test_proximal_freeb_jac_blocked        |   -0.03 %    |     7.883e+03      |     7.880e+03      |    -2.54     |       66.60        |       67.12        |
  test_proximal_freeb_jac_batched        |    0.36 %    |     7.840e+03      |     7.868e+03      |    27.95     |       65.42        |       66.23        |
  test_proximal_jac_ripple               |    0.63 %    |     3.709e+03      |     3.732e+03      |    23.34     |       42.09        |       42.56        |
  test_proximal_jac_ripple_bounce1d      |   -0.11 %    |     3.965e+03      |     3.961e+03      |    -4.52     |       52.53        |       54.12        |
  test_eq_solve                          |   -8.74 %    |     1.999e+03      |     1.824e+03      |   -174.78    |       38.14        |       38.73        |
  test_objective_quadratic_flux_jac      |    0.30 %    |     2.707e+03      |     2.715e+03      |     8.00     |       43.77        |       44.35        |

For the memory plots, go to the summary of Memory Benchmarks workflow and download the artifact.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 94.37%. Comparing base (fcc29be) to head (336c049).
⚠️ Report is 1 commits behind head on master.

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #2281   +/-   ##
=======================================
  Coverage   94.37%   94.37%           
=======================================
  Files         101      101           
  Lines       28936    28953   +17     
=======================================
+ Hits        27308    27325   +17     
  Misses       1628     1628           

Files with missing lines	Coverage Δ
desc/optimize/aug_lagrangian_ls.py	96.00% <100.00%> (+0.11%)	⬆️
desc/optimize/fmin_scalar.py	98.34% <100.00%> (+<0.01%)	⬆️
desc/optimize/least_squares.py	99.45% <100.00%> (+0.01%)	⬆️
desc/optimize/utils.py	95.43% <100.00%> (+0.05%)	⬆️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[YigitElma]

For historical comparison, here are the same profiles from #1688 (comment)

• Current master vs this PR (jac_chunk_size=500)

• Current master vs this PR (jac_chunk_size=1)

Note: I was sampling memory at around 7kHz; this should resolve almost all peaks.

SVD/Cholesky

• SVD (jac_chunk_size=500)

• Cholesky (jac_chunk_size=500)

[github-actions]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     +3.44 +/- 2.58     | +2.81e-02 +/- 2.10e-02 |  8.45e-01 +/- 1.7e-02  |  8.16e-01 +/- 1.2e-02  |
 test_equilibrium_init_lowres            |     +1.33 +/- 3.51     | +8.50e-02 +/- 2.24e-01 |  6.46e+00 +/- 1.7e-01  |  6.38e+00 +/- 1.5e-01  |
 test_objective_compile_atf              |     -0.52 +/- 4.04     | -3.03e-02 +/- 2.35e-01 |  5.79e+00 +/- 1.3e-01  |  5.82e+00 +/- 1.9e-01  |
 test_objective_compute_atf              |     +1.26 +/- 7.31     | +2.60e-05 +/- 1.51e-04 |  2.09e-03 +/- 1.4e-04  |  2.07e-03 +/- 4.7e-05  |
 test_objective_jac_atf                  |     +3.58 +/- 5.30     | +5.57e-02 +/- 8.24e-02 |  1.61e+00 +/- 7.1e-02  |  1.55e+00 +/- 4.1e-02  |
 test_perturb_1                          |     +1.87 +/- 2.36     | +2.12e-01 +/- 2.68e-01 |  1.16e+01 +/- 2.6e-01  |  1.13e+01 +/- 7.8e-02  |
 test_proximal_jac_atf                   |     +0.69 +/- 1.75     | +3.63e-02 +/- 9.21e-02 |  5.30e+00 +/- 6.9e-02  |  5.26e+00 +/- 6.1e-02  |
 test_proximal_freeb_compute             |     +0.90 +/- 3.66     | +1.22e-03 +/- 4.93e-03 |  1.36e-01 +/- 2.8e-03  |  1.35e-01 +/- 4.1e-03  |
 test_solve_fixed_iter                   |     +0.07 +/- 2.04     | +1.72e-02 +/- 4.79e-01 |  2.36e+01 +/- 3.5e-01  |  2.35e+01 +/- 3.3e-01  |
 test_LinearConstraintProjection_build   |     +0.59 +/- 4.76     | +3.86e-02 +/- 3.13e-01 |  6.61e+00 +/- 1.2e-01  |  6.57e+00 +/- 2.9e-01  |
 test_objective_compute_ripple           |     -1.46 +/- 4.55     | -3.08e-03 +/- 9.61e-03 |  2.08e-01 +/- 5.1e-03  |  2.11e-01 +/- 8.2e-03  |
 test_objective_grad_ripple              |     +1.25 +/- 3.03     | +1.12e-02 +/- 2.70e-02 |  9.04e-01 +/- 1.5e-02  |  8.93e-01 +/- 2.2e-02  |
 test_objective_quadratic_flux_compute   |     +0.31 +/- 13.71    | +1.69e-04 +/- 7.39e-03 |  5.41e-02 +/- 2.8e-03  |  5.39e-02 +/- 6.8e-03  |
 test_build_transform_fft_midres         |     -0.91 +/- 3.78     | -8.06e-03 +/- 3.37e-02 |  8.82e-01 +/- 2.2e-02  |  8.90e-01 +/- 2.6e-02  |
 test_build_transform_fft_highres        |     -0.78 +/- 2.83     | -9.10e-03 +/- 3.32e-02 |  1.16e+00 +/- 2.8e-02  |  1.17e+00 +/- 1.8e-02  |
 test_equilibrium_init_medres            |     +2.24 +/- 4.02     | +1.54e-01 +/- 2.76e-01 |  7.02e+00 +/- 2.0e-01  |  6.86e+00 +/- 1.9e-01  |
 test_objective_compile_dshape_current   |     +0.85 +/- 1.62     | +3.32e-02 +/- 6.30e-02 |  3.93e+00 +/- 4.9e-02  |  3.89e+00 +/- 3.9e-02  |
 test_objective_compute_dshape_current   |     +2.04 +/- 12.04    | +1.37e-05 +/- 8.08e-05 |  6.85e-04 +/- 4.1e-05  |  6.71e-04 +/- 6.9e-05  |
 test_objective_jac_dshape_current       |     +0.02 +/- 28.29    | +3.79e-06 +/- 6.92e-03 |  2.45e-02 +/- 4.9e-03  |  2.44e-02 +/- 4.9e-03  |
 test_perturb_2                          |     +2.28 +/- 2.62     | +3.43e-01 +/- 3.94e-01 |  1.54e+01 +/- 2.4e-01  |  1.50e+01 +/- 3.1e-01  |
 test_proximal_jac_atf_with_eq_update    |     -0.07 +/- 0.93     | -8.36e-03 +/- 1.15e-01 |  1.23e+01 +/- 9.2e-02  |  1.23e+01 +/- 6.8e-02  |
 test_proximal_freeb_jac                 |     -1.08 +/- 4.18     | -5.04e-02 +/- 1.94e-01 |  4.60e+00 +/- 1.2e-01  |  4.65e+00 +/- 1.5e-01  |
 test_solve_fixed_iter_compiled          |     -1.78 +/- 2.17     | -1.09e-01 +/- 1.33e-01 |  6.03e+00 +/- 1.1e-01  |  6.14e+00 +/- 8.2e-02  |
 test_objective_compute_ripple_bounce1d  |     +0.74 +/- 3.91     | +2.05e-03 +/- 1.09e-02 |  2.80e-01 +/- 6.5e-03  |  2.78e-01 +/- 8.7e-03  |
 test_objective_grad_ripple_bounce1d     |     +0.28 +/- 2.09     | +2.67e-03 +/- 2.03e-02 |  9.72e-01 +/- 1.2e-02  |  9.69e-01 +/- 1.6e-02  |
 test_objective_quadratic_flux_jac       |     +0.28 +/- 0.59     | +2.35e-02 +/- 5.00e-02 |  8.52e+00 +/- 3.5e-02  |  8.49e+00 +/- 3.6e-02  |

Github CI performance can be noisy. When evaluating the benchmarks, developers should take this into account.