Use a custom train_step for numerical stability

[kmdalton]

Modify careless to

• track the norm of the gradients
• cleanly terminate optimization on numerical issues
• possibly mitigate some numerical issues

[codecov-commenter]

Codecov Report

Merging #126 (c7796fe) into main (7b36c94) will decrease coverage by 0.42%. Report is 3 commits behind head on main. The diff coverage is 20.00%.

@@            Coverage Diff             @@
##             main     #126      +/-   ##
==========================================
- Coverage   78.92%   78.51%   -0.42%     
==========================================
  Files          50       50              
  Lines        2164     2220      +56     
==========================================
+ Hits         1708     1743      +35     
- Misses        456      477      +21     

Flag	Coverage Δ
unittests	78.51% <20.00%> (-0.42%)	:arrow_down:

Flags with carried forward coverage won't be shown. Click here to find out more.

Files Changed	Coverage Δ
careless/models/merging/variational.py	71.42% <20.00%> (-14.16%)	:arrow_down:

... and 4 files with indirect coverage changes

:mega: We’re building smart automated test selection to slash your CI/CD build times. Learn more