Better Steady State Adjoint

[avik-pal]

• External Heuristics Control
  • NoHeuristic --> Follow exactly what user wants
  • SmallJACHeuristic --> If Jacobian is small then construct it ignoring user requirements
• concrete_jac -- Force Jacobian Computation
• Assume Uniform Square Block Diagonal (is there a mathematical name for these matrices where all the blocks are of uniform size?) -- Most Machine Learning Applications. But we can't default to it since the user might use BatchNorm which messes up the sparsity pattern assumption.
  • If concrete jac -- use sparsity pattern to collapse Jacobian computation of an entire batch to a single batch (should scale really well)
  • Use a specialized GMRES solver for this part.
• Supports arbitrary parameter structures

[ai-maintainer]

AI-Maintainer Review for PR - Better Steady State Adjoint

Title and Description ⚠️

Title and Description Need More Clarity

The title of the pull request suggests improvements to the Steady State Adjoint functionality, but the description does not provide a clear explanation of the purpose of the changes. It would be helpful to provide a more detailed and concise description of the purpose of the changes in the pull request.

Scope of Changes 👍

Changes are Narrowly Focused

The changes are focused on improving the Steady State Adjoint functionality. The modifications to several files related to the Steady State Adjoint algorithm, sensitivity functions, and problem setup all seem to be related to the same topic of enhancing the Steady State Adjoint.

Documentation ⚠️

Missing Docstrings

The following functions, classes, or methods need docstrings to describe their behavior, arguments, and return values:

• SteadyStateAdjoint{CJ, BD, CS, AD, FDT, VJP, LS} struct
• setvjp function
• SteadyStateAdjointSensitivityFunction function
• SteadyStateAdjointProblem function

Testing ⚠️

No Information on Testing

The description does not mention how the changes were tested. It would be beneficial for the author to include details about the testing methodology, such as specific test cases, frameworks, or scenarios used to validate the changes.

Suggested Changes

• Please provide a more detailed and concise description of the purpose of the changes in the pull request.
• Add docstrings to the SteadyStateAdjoint{CJ, BD, CS, AD, FDT, VJP, LS} struct, setvjp function, SteadyStateAdjointSensitivityFunction function, and SteadyStateAdjointProblem function to describe their behavior, arguments, and return values.
• Include details about the testing methodology, such as specific test cases, frameworks, or scenarios used to validate the changes.

Reviewed with AI Maintainer

[avik-pal]

Needs SciML/LinearSolve.jl#366 and a tag of SparseDiffTools.jl

[ChrisRackauckas]

was this supposed to pass tests?

[avik-pal]

was this supposed to pass tests?

Yup 😓. If it doesn't pass now I will have to debug furthur

[codecov]

Codecov Report

Merging #877 (28cfa9c) into master (da8b6d7) will decrease coverage by 0.36%.
The diff coverage is 67.17%.

❗ Current head 28cfa9c differs from pull request most recent head d7c4635. Consider uploading reports for the commit d7c4635 to get more accurate results

@@            Coverage Diff             @@
##           master     #877      +/-   ##
==========================================
- Coverage   61.91%   61.56%   -0.36%     
==========================================
  Files          20       20              
  Lines        4385     4472      +87     
==========================================
+ Hits         2715     2753      +38     
- Misses       1670     1719      +49     

Files	Coverage Δ
src/SciMLSensitivity.jl	100.00% <ø> (ø)
src/adjoint_common.jl	94.42% <100.00%> (-0.28%)	⬇️
src/derivative_wrappers.jl	89.54% <57.14%> (-0.66%)	⬇️
src/sensitivity_algorithms.jl	80.95% <85.71%> (+3.17%)	⬆️
src/steadystate_adjoint.jl	84.81% <82.50%> (-3.33%)	⬇️
src/concrete_solve.jl	66.57% <45.28%> (-2.82%)	⬇️

... and 5 files with indirect coverage changes

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

[ChrisRackauckas]

It looks like this causes a regression in the vjps? See the Core 3 test failures.

[avik-pal]

Closing this but not deleting the branch.