Work toward parallelization via sharding (multi-device) composed with vmap (across memory in single device)

[unalmis]

Pulls just the code from ku/shard on #1773 that other developers said they were okay with merging on that day.

[github-actions]

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |    4.49 %    |     3.841e+03      |     4.014e+03      |    172.49    |       33.86        |       31.64        |
  test_proximal_jac_w7x_with_eq_update   |    0.49 %    |     6.690e+03      |     6.723e+03      |    32.63     |       161.33       |       158.41       |
  test_proximal_freeb_jac                |   -0.28 %    |     1.320e+04      |     1.316e+04      |    -36.82    |       78.04        |       78.03        |
  test_proximal_freeb_jac_blocked        |   -1.19 %    |     7.657e+03      |     7.566e+03      |    -91.19    |       69.49        |       68.90        |
  test_proximal_freeb_jac_batched        |    0.68 %    |     7.528e+03      |     7.579e+03      |    51.35     |       70.11        |       69.18        |
  test_proximal_jac_ripple               |    0.68 %    |     7.548e+03      |     7.599e+03      |    51.21     |       69.31        |       69.69        |
  test_proximal_jac_ripple_spline        |   -0.20 %    |     3.518e+03      |     3.511e+03      |    -7.21     |       72.32        |       72.82        |
  test_eq_solve                          |   -0.50 %    |     2.072e+03      |     2.062e+03      |    -10.30    |       125.78       |       125.99       |

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

[codecov]

Codecov Report

❌ Patch coverage is 85.36585% with 6 lines in your changes missing coverage. Please review.
✅ Project coverage is 95.81%. Comparing base (e0dcf9e) to head (febfbb0).
⚠️ Report is 61 commits behind head on master.

Files with missing lines	Patch %	Lines
desc/batching.py	85.36%	6 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1869      +/-   ##
==========================================
- Coverage   95.83%   95.81%   -0.02%     
==========================================
  Files         100      100              
  Lines       27511    27535      +24     
==========================================
+ Hits        26365    26383      +18     
- Misses       1146     1152       +6     

Files with missing lines	Coverage Δ
desc/batching.py	88.35% <85.36%> (-1.95%)	⬇️

... 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.

[github-actions]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     +2.54 +/- 5.86     | +1.43e-02 +/- 3.30e-02 |  5.77e-01 +/- 3.3e-02  |  5.63e-01 +/- 5.4e-03  |
 test_equilibrium_init_medres            |     +1.49 +/- 4.08     | +7.43e-02 +/- 2.03e-01 |  5.05e+00 +/- 1.6e-01  |  4.98e+00 +/- 1.2e-01  |
 test_equilibrium_init_highres           |     +1.08 +/- 2.91     | +6.00e-02 +/- 1.62e-01 |  5.63e+00 +/- 1.4e-01  |  5.57e+00 +/- 8.1e-02  |
 test_objective_compile_dshape_current   |     +1.26 +/- 3.18     | +4.21e-02 +/- 1.06e-01 |  3.38e+00 +/- 1.0e-01  |  3.34e+00 +/- 3.4e-02  |
 test_objective_compute_dshape_current   |     +3.36 +/- 8.22     | +2.47e-05 +/- 6.03e-05 |  7.59e-04 +/- 5.3e-05  |  7.34e-04 +/- 2.9e-05  |
 test_objective_jac_dshape_current       |     -1.21 +/- 17.00    | -3.93e-04 +/- 5.53e-03 |  3.21e-02 +/- 3.8e-03  |  3.25e-02 +/- 4.0e-03  |
 test_perturb_2                          |     +0.36 +/- 3.60     | +6.05e-02 +/- 6.12e-01 |  1.70e+01 +/- 4.2e-01  |  1.70e+01 +/- 4.4e-01  |
 test_proximal_jac_atf_with_eq_update    |     +0.10 +/- 0.76     | +1.35e-02 +/- 1.04e-01 |  1.38e+01 +/- 7.4e-02  |  1.38e+01 +/- 7.3e-02  |
 test_proximal_freeb_jac                 |     +0.65 +/- 6.26     | +3.31e-02 +/- 3.17e-01 |  5.11e+00 +/- 3.2e-01  |  5.07e+00 +/- 3.6e-02  |
 test_solve_fixed_iter_compiled          |     +0.41 +/- 1.11     | +6.91e-02 +/- 1.89e-01 |  1.71e+01 +/- 1.6e-01  |  1.70e+01 +/- 9.4e-02  |
 test_LinearConstraintProjection_build   |     -0.26 +/- 2.24     | -2.20e-02 +/- 1.90e-01 |  8.42e+00 +/- 1.5e-01  |  8.45e+00 +/- 1.2e-01  |
 test_objective_compute_ripple_spline    |     -0.19 +/- 2.53     | -5.63e-04 +/- 7.36e-03 |  2.91e-01 +/- 5.1e-03  |  2.91e-01 +/- 5.3e-03  |
 test_objective_grad_ripple_spline       |     -0.19 +/- 2.62     | -2.14e-03 +/- 2.95e-02 |  1.12e+00 +/- 2.2e-02  |  1.12e+00 +/- 2.0e-02  |
 test_build_transform_fft_midres         |     -0.03 +/- 6.80     | -2.14e-04 +/- 4.75e-02 |  6.99e-01 +/- 4.5e-02  |  6.99e-01 +/- 1.6e-02  |
 test_build_transform_fft_highres        |     -0.34 +/- 3.58     | -3.12e-03 +/- 3.31e-02 |  9.19e-01 +/- 2.3e-02  |  9.23e-01 +/- 2.3e-02  |
 test_equilibrium_init_lowres            |     -4.48 +/- 4.87     | -2.04e-01 +/- 2.22e-01 |  4.36e+00 +/- 9.1e-02  |  4.56e+00 +/- 2.0e-01  |
 test_objective_compile_atf              |     -0.17 +/- 5.42     | -1.03e-02 +/- 3.34e-01 |  6.15e+00 +/- 1.9e-01  |  6.16e+00 +/- 2.7e-01  |
 test_objective_compute_atf              |     +9.24 +/- 42.16    | +1.81e-04 +/- 8.28e-04 |  2.14e-03 +/- 4.5e-04  |  1.96e-03 +/- 7.0e-04  |
 test_objective_jac_atf                  |     -2.00 +/- 2.26     | -3.48e-02 +/- 3.92e-02 |  1.70e+00 +/- 3.6e-02  |  1.74e+00 +/- 1.6e-02  |
 test_perturb_1                          |     -1.58 +/- 4.38     | -2.19e-01 +/- 6.06e-01 |  1.36e+01 +/- 3.0e-01  |  1.39e+01 +/- 5.2e-01  |
 test_proximal_jac_atf                   |     +0.12 +/- 1.12     | +6.74e-03 +/- 6.20e-02 |  5.55e+00 +/- 2.4e-02  |  5.54e+00 +/- 5.7e-02  |
 test_proximal_freeb_compute             |     -0.86 +/- 2.92     | -1.38e-03 +/- 4.69e-03 |  1.59e-01 +/- 2.5e-03  |  1.60e-01 +/- 4.0e-03  |
 test_solve_fixed_iter                   |     -2.40 +/- 1.57     | -7.07e-01 +/- 4.63e-01 |  2.88e+01 +/- 4.4e-01  |  2.95e+01 +/- 1.4e-01  |
 test_objective_compute_ripple           |     +0.15 +/- 0.97     | +3.99e-03 +/- 2.53e-02 |  2.62e+00 +/- 2.0e-02  |  2.62e+00 +/- 1.5e-02  |
 test_objective_grad_ripple              |     +0.20 +/- 2.01     | +9.53e-03 +/- 9.47e-02 |  4.73e+00 +/- 2.8e-02  |  4.72e+00 +/- 9.1e-02  |

[unalmis]

Yigit can you override codecov for merge? Thank you

[YigitElma]

@f0uriest @dpanici, you can merge by override when you approve

[rahulgaur104]

I suspect this to be the reason behind #1925

[dpanici]

This is on us, the jax tests all failed. I did not notice that