Ku/einsum

[unalmis]

use einsum and the like to save memory since at this point it's clear jax won't fix the memory regression in 0.5.3

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[github-actions]

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |    3.81 %    |     3.859e+03      |     4.006e+03      |    147.13    |       40.37        |       36.36        |
  test_proximal_jac_w7x_with_eq_update   |   -2.68 %    |     6.594e+03      |     6.417e+03      |   -176.73    |       162.09       |       160.98       |
  test_proximal_freeb_jac                |   -0.01 %    |     1.320e+04      |     1.320e+04      |    -1.61     |       86.14        |       82.11        |
  test_proximal_freeb_jac_blocked        |    0.31 %    |     7.493e+03      |     7.517e+03      |    23.51     |       73.20        |       72.51        |
  test_proximal_freeb_jac_batched        |   -0.89 %    |     7.487e+03      |     7.420e+03      |    -66.64    |       73.01        |       72.28        |
  test_proximal_jac_ripple               |    0.41 %    |     3.484e+03      |     3.498e+03      |    14.18     |       64.00        |       64.89        |
  test_proximal_jac_ripple_bounce1d      |    0.33 %    |     3.518e+03      |     3.530e+03      |    11.63     |       85.40        |       75.43        |
  test_eq_solve                          |    0.70 %    |     1.988e+03      |     2.002e+03      |    13.99     |       93.77        |       93.06        |

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 95.44%. Comparing base (2715006) to head (be42f20).
⚠️ Report is 1 commits behind head on ku/bounce.

Additional details and impacted files

@@              Coverage Diff              @@
##           ku/bounce    #2033      +/-   ##
=============================================
+ Coverage      95.43%   95.44%   +0.01%     
=============================================
  Files            103      103              
  Lines          28451    28451              
=============================================
+ Hits           27153    27156       +3     
+ Misses          1298     1295       -3     

Files with missing lines	Coverage Δ
desc/integrals/_bounce_utils.py	92.60% <100.00%> (ø)
desc/integrals/_interp_utils.py	97.42% <100.00%> (ø)
desc/integrals/basis.py	85.83% <ø> (ø)
desc/objectives/_fast_ion.py	95.00% <ø> (ø)
desc/objectives/_neoclassical.py	94.80% <ø> (ø)

... and 3 files with indirect coverage changes

🚀 New features to boost your workflow:

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