Stanc.Std: Move over Analysis module

[WardBrian]

The most notable change here is the addition of our own vendored polymorphic set. This is based on OCaml's set implementation, but with polymorphic comparison swapped in. Remarkably, I've found that it is dramatically faster than Base.Set.Poly in our uses, presumably because Base.Set.Poly is storing the comparison function in the runtime representation still, even though it never changes

Submission Checklist

• Run unit tests
• Documentation
  • If a user-facing facing change was made, the documentation PR is here:
  • OR, no user-facing changes were made

Release notes

Copyright and Licensing

By submitting this pull request, the copyright holder is agreeing to
license the submitted work under the BSD 3-clause license (https://opensource.org/licenses/BSD-3-Clause)

---

_{Stack created with GitHub Stacks CLI • Give Feedback 💬}

[codecov]

Codecov Report

❌ Patch coverage is 93.33333% with 28 lines in your changes missing coverage. Please review.
✅ Project coverage is 92.11%. Comparing base (38752c2) to head (fcbb41b).

Files with missing lines	Patch %	Lines
src/analysis_and_optimization/Factor_graph.ml	78.26%	10 Missing ⚠️
src/analysis_and_optimization/Memory_patterns.ml	90.90%	4 Missing ⚠️
src/analysis_and_optimization/Mir_utils.ml	81.81%	4 Missing ⚠️
...rc/analysis_and_optimization/Monotone_framework.ml	97.58%	3 Missing ⚠️
src/analysis_and_optimization/Optimize.ml	94.82%	3 Missing ⚠️
src/analysis_and_optimization/Pedantic_analysis.ml	95.91%	2 Missing ⚠️
src/analysis_and_optimization/Partial_evaluator.ml	83.33%	1 Missing ⚠️
src/std/std.ml	85.71%	1 Missing ⚠️

Additional details and impacted files

@@                    Coverage Diff                    @@
##           std/stan_math_backend    #1661      +/-   ##
=========================================================
- Coverage                  92.16%   92.11%   -0.05%     
=========================================================
  Files                         70       70              
  Lines                      10126    10139      +13     
=========================================================
+ Hits                        9333     9340       +7     
- Misses                       793      799       +6     

Files with missing lines	Coverage Δ
src/analysis_and_optimization/Dataflow_types.ml	25.00% <100.00%> (+25.00%)	⬆️
src/analysis_and_optimization/Dataflow_utils.ml	100.00% <100.00%> (ø)
...analysis_and_optimization/Debug_data_generation.ml	81.63% <100.00%> (ø)
...c/analysis_and_optimization/Dependence_analysis.ml	100.00% <100.00%> (ø)
...nalysis_and_optimization/Pedantic_dist_warnings.ml	91.70% <100.00%> (ø)
src/middle/Index.ml	85.18% <ø> (ø)
src/stan_math_backend/Transform_Mir.ml	95.54% <100.00%> (ø)
src/analysis_and_optimization/Partial_evaluator.ml	90.95% <83.33%> (+0.02%)	⬆️
src/std/std.ml	90.82% <85.71%> (+0.53%)	⬆️
src/analysis_and_optimization/Pedantic_analysis.ml	96.08% <95.91%> (+0.49%)	⬆️
... and 5 more

... and 2 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[nhuurre]

Looks like constant propagation got a bit smarter, but why?

[WardBrian]

I spent a few hours on trying to debug this and I couldn’t pinpoint the change.

I also threw whatever free credits of Claude github gives me, and it spun in circles but couldn’t identify anything either after spending the entire month’s allocation.

In the end, I think the fact that this wasn’t being propagated in master was a bug, and the relevant optimizations are only turned on for Oexperimental, so I gave up looking for the cause… if you can find anything without wasting too much time, I would still be curious!

[nhuurre]

Yes, this was probably a bug in master and not worth chasing down. However, note that the anomaly happens in constant propagation, which is enabled for O1, even though it didn't trigger without some experimental passes first. If you're curious and want to waste even more time I found a "minimal" example:

let%expect_test "constant propagation" =
  let mir =
    Test_utils.mir_of_string {|
data {
  vector[2] X;
  int sum_y;
  real alpha_occ;
  vector[2] logit_p;
}
generated quantities {
  vector[2] logit_psi = alpha_occ + alpha_occ * X;
  array[2] real psi_con;
  if (1) {
    real psi = inv_logit(logit_psi[1]);
    real qT = prod(logit_p);
    psi_con[1] = inv(psi * qT + 1);
  } else {
    psi_con[1] = 1;
  }
  real psi = inv_logit(logit_psi[2]);
  real qT = prod(logit_p);
  psi_con[2] = inv(psi * qT + 1);
}     |} in
  let open Analysis_and_optimization.Optimize in
  let mir = expression_propagation mir in
  let mir = lazy_code_motion mir in
  let mir = constant_propagation mir in
  Fmt.str "@[<v>%a@]" Middle.Program.Typed.pp mir |> print_endline;
  [%expect {| ... |}]

[WardBrian]

The if (1) seems to be the load-bearing part. If you change it to something semantically equivalent, like if (2), then the constant does indeed get moved down

[WardBrian]

My theories:

• Some part of the optimization pipeline is iteration-order-dependent by accident, and this branch gets a favorable order
• There is something in master that is accidentally sensitive to the location metadata, and isn't here (though I've tried to find the likely places this would be, and came up empty)

Assuming we're both happy not knowing, what do you think of the rest of the changes?

[nhuurre]

Accidental order-dependence sounds plausible. The rest of the changes look good.