Initial implicit/truncated backward modes

[bamos]

Migrating this from the private repo (where we have some more context. I've updated the example/test to use the quadratic cost fitting example as I like it's simplicity to the older example we were using with the cost weights. Here's what the current backward pass mode example looks like:

--- backward_mode=FULL
[ 0.09670612 -0.7900081   0.08778726  0.13497931  0.21781093 -0.20041543
  0.11018689 -1.5431342   0.13462512 -0.23546126]

--- backward_mode=IMPLICIT
[ 0.04835452 -0.3950162   0.04389496  0.0674917   0.10890876 -0.10021076
  0.05509511 -0.77159065  0.0673146  -0.11773425]

--- backward_mode=TRUNCATED, backward_num_iterations=5
[ 0.09368691 -0.7653437   0.08504649  0.1307652   0.21101077 -0.19415839
  0.10674679 -1.4949567   0.13042209 -0.22811009]

--- Numeric derivative
[ 0.09670693 -0.7900125   0.08778608  0.13497717  0.21781142 -0.20041636
  0.11019139 -1.54313637  0.13462317 -0.2354647 ]

=== Runtimes
Forward: 1.60e-02 s +/- 1.57e-03 s
Backward (FULL): 4.61e-03 s +/- 3.73e-04 s
Backward (IMPLICIT) 5.18e-04 s +/- 7.88e-05 s
Backward (TRUNCATED, 5 steps) 1.59e-03 s +/- 8.15e-05 s

This is almost ready to merge but still WIP as I'm still debugging some scaling issues in the implicit mode, which should be easier to debug with this simpler example.

[mhmukadam]

Looks great! Left some minor comments, after which we should merge this and try to debug the implicit grad offset issue in a separate PR (unless you already fixed it).

Something failed in the GPU test. Might be a CI issue, don't think it was related to any commits here.

[bamos]

Hi @mhmukadam and @luisenp, I just went through and addressed all the inline comments. Let me know if there's anything else! Otherwise this is ready to merge from my perspective, and I'll keep debugging some of the gradient scaling issues separately for #39.

[luisenp]

LGTM, some minor comments and a question!

[mhmukadam]

Looks good!

[bamos]

Thanks for the pass @luisenp! I just addressed your comments and pushed the latest that fixes the scaling issues and passes the backward tests with atol=1e-4. Can you let me know if there's anything else?