[avehtari]

Replace the incomplete beta gradient roots with a continued-fraction implementation from TensorFlow Probability

Big increase in stability, accuracy and speed.
Fixes #2417.

This PR was Claude assisted

Summary

grad_reg_inc_beta, inc_beta_dda and inc_beta_ddb compute the derivatives of the regularized incomplete beta function I_z(a, b) with respect to the shapes a and b. develop has two separate implementations of these derivatives. Path A (inc_beta_dda, inc_beta_ddb) serves beta_cdf, neg_binomial_cdf and neg_binomial_2_cdf. Path B (grad_reg_inc_beta, through grad_2F1) serves the lcdf and lccdf of beta, beta_proportion, neg_binomial and student_t, and also student_t_cdf and the autodiff of inc_beta. Both paths return NaN gradients, or gradients wrong by O(1), at large shapes, at b ≫ a, and where I is near 1.

This PR replaces both paths with one implementation: the Boik and Robinson-Cox (1998) algorithm as TensorFlow Probability implements it. It uses a power series near the endpoints and a modified-Lentz continued fraction elsewhere. It differentiates both term by term and evaluates the prefactor in log space. inc_beta_dda and inc_beta_ddb become wrappers, and beta_cdf calls the root once instead of twice. No signature changes.

The values of the distributions do not change. Only the gradients change.

I compiled both trees and ran them on the same 18 000 points in nine regions (details below):

	develop	branch
non-finite, path B, large shapes near the mode	1325 of 2000	0
non-finite, path A (throws), all regions	100	0
gradient of log I wrong by > 1e-2, path A, b ≫ a	408 of 2000	0
gradient of log I wrong by > 1e-2, path B, negbin n > 200	1203 of 1999	0
worst error of the gradient of log I, branch	n/a	1e-11
worst relative error where the derivative is not negligible, branch	n/a	3.7e-9
spread between double, var, fvar, fvar<fvar>	0	0
speed, grad_reg_inc_beta in double	1x	30x
speed, beta_lcdf per observation, α = β = 5	1x	40x
speed, beta_cdf per observation, α = β = 50	1x	2.4x

The defects in the current code on develop

This section describes develop at 5252d51d47. The reference is mpmath at 80 digits, computed by two independent routes: quadrature of the density, and numerical differentiation at 110 digits. The two routes agree to 60 digits at 54 check points. All develop numbers come from the compiled develop headers.

1. grad_reg_inc_beta returns NaN for every parameter when beta(a, b) underflows. Its callers pass beta(a, b) as betaAB, and the function divides by it. beta(a, b) underflows when a is near b and a + b is above about 1100. For example, beta_lcdf(y | 600, 600) has a NaN gradient for both shapes, and so does inc_beta(a, b, z) with autodiff shapes. In the sweep region with a, b in [100, 12500] and z near the mode, where I is of order 1, 1325 of 2000 points are NaN.

2. grad_reg_inc_beta drops the hypergeometric term when z^a (1-z)^b underflows. When C = z^a (1-z)^b / a is 0, the function skips grad_2F1 and returns I (log z - 1/a - ψ(a) + ψ(a+b)). When I is near 1, the dropped term is the whole answer, because 2F1(a+b, 1; a+1; z) grows like (1-z)^(-b). beta_lcdf(0.758 | 352, 590) returns d/dα = 0.705 for a true -3.4e-138. neg_binomial_lcdf(132 | 313, 0.068) returns d/dα = 0 for a true -2.3.

3. inc_beta_dda and inc_beta_ddb stop on an absolute threshold that can be met before the first term. Both stop the series when fabs(summand) < 1e-10. The first summand contains the factor ((a+1)/(a+b))^3, which is below 1e-10 for b ≫ a. The loop then exits at once and returns the first term. beta_cdf(1.8e-4 | 1.145, 6786) returns d/dα = -0.203 for a true -0.386, at I = 0.65. In the sweep region with b ≫ a and small z, the gradient of log I has an error above 1e-6 at 1301 of 2000 points, and above 1e-2 at 408.

4. inc_beta_dda and inc_beta_ddb reflect on fixed z thresholds, not on which side is small. For z > 0.75, and in other fixed bands, they evaluate 1 - I_{1-z}(b, a) wherever the mass is. In the deep lower tail the result is then I' ≈ 1 times a bracket that must cancel to 1e-47, and the computed bracket is rounding noise of size 1e-15. beta_cdf(0.76 | 446, 5) returns d/dα = 8.9e-16 for a true -1.06e-47, so the gradient of log(beta_cdf(...)) is 2.5e+31. At some large-shape points the loop also reaches the k > 1e5 iteration guard and throws, inside the range that the docstring calls tested.

5. The existing unit tests assert wrong values. inc_beta_ddb_test.cpp passes digamma(a) as the fourth argument, which is digamma(b), and its expected constants come from that call (reported in #2417). inc_beta_dda_test.cpp and inc_beta_ddb_test.cpp also assert the output of develop at points where it is wrong by orders of magnitude: 1.8226241 for a true -3.94e-5, and 9.3959293 for a true 1.48e-8. grad_reg_inc_beta_test.cpp expects d/db = -inf at z = 1, where I ≡ 1 and the derivative is 0.

What this PR does

The new body of stan/math/prim/fun/grad_reg_inc_beta.hpp is a C++ port of _betainc_partials from TensorFlow Probability (tensorflow_probability/python/math/special.py, Apache-2.0; see the licence note below). The algorithm is from Boik and Robinson-Cox (1998), Derivatives of the Incomplete Beta Function, JSS 3(1). It evaluates the DLMF 8.17.22 continued fraction with the modified Lentz method and carries the derivatives of the partial numerators through the recurrence. Near the endpoints it uses the power series from Cephes, also differentiated term by term.

Each defect above has a direct repair. The port evaluates the prefactor z^a (1-z)^b / (a B(a, b)) in log space, which removes defects 1 and 2. It applies the symmetry relation I_z(a, b) = 1 - I_{1-z}(b, a) at the mode for the continued fraction and at the mean for the series, which removes defect 4. It has no absolute threshold: the series stops at eps / a and the continued fraction at 3 eps, which removes defect 3.

The port changes the two stop tests of the TFP source, one in the series and one in the continued fraction. The code marks both changes. TFP stops when the value has converged. When b is a positive integer, the series factor (n - b) and the continued-fraction numerator d_{2m} are exactly 0 at n = m = b. The value has then converged, but the derivatives still need the rest of the series. TFP's own d/db is off by 4.2e-4 at (2, 3, 0.25) and by 6.7 % at (10, 1, 0.5). Its d/da is off by 5.6 % at (1, 100, 0.01), where the symmetry relation makes b' = 1. The port stops only when the derivative increments have also converged. Integer shapes are common, for example as initial values or as fixed shapes. I will report the TFP defect upstream.

inc_beta_dda and inc_beta_ddb become wrappers over grad_reg_inc_beta that return one of the two derivatives. beta_cdf calls grad_reg_inc_beta once, not inc_beta_dda and inc_beta_ddb separately. prim/fun.hpp now includes inc_beta_dda.hpp, inc_beta_ddb.hpp and inc_beta_ddz.hpp; before, it reached them only through prob/beta_cdf.hpp. grad_reg_inc_beta keeps its betaAB argument for interface compatibility but no longer uses it.

Net change: about 180 lines of algorithm replace 123.

Accuracy after the change

Same 18 000 points, same reference, both trees compiled with the same probe. The metric is the absolute error of the gradient of log I. beta_lcdf returns this gradient, and autodiff of log(beta_cdf(...)) gives it too. A relative error is not useful for a derivative that is 1e-300 of I.

region (2 000 points each)	develop: non-finite / error > 1e-2	branch: worst error
a, b < 1	0 / 0	1e-13
1 ≤ a, b ≤ 20, z < 0.75	0 / 0	3e-14
z > 0.75, a < 500	path A 3 / 81, path B 3 / 125	7e-15
b > a, b > 500, z ≤ 0.75	path A 10 / 86, path B 550 / 933	3e-13 (1)
b ≫ a, z small	path A 2 / 408, path B 2 / 12	3.5e-9
large a, b, z uniform	path A 46 / 36, path B 932 / 235	4e-12 (1)
large a, b, z near the mode	path A 0 / 0, path B 1325 / 9	1e-13
negbin pattern, n ≤ 200	path A 0 / 4, path B 0 / 26	4e-13
negbin pattern, 200 < n ≤ 5000	path A 39 / 22, path B 292 / 1203	4e-13

(1) Excluding points where I itself is subnormal in double.

Where the derivative is not negligible (|dI| / I ≥ 1e-3), the relative error is at most 6e-13 in every region except b ≫ a with small z. There, one continued-fraction point at the symmetry switch reaches 3.7e-9. The four scalar types double, var, fvar<double> and fvar<fvar<double>> give bit-identical values on both trees. On the branch, the autodiff types have the same error profile as double.

Speed

Xeon E5-2680 v3 at 2.50 GHz, exclusive node, performance governor, two rounds that agree within 2 %. Grid: a, b in [0.5, 60], z in [0.05, 0.95]. The benchmark loop adds a term that depends on the accumulator to every input, so the compiler cannot move the kernel out of the loop. I subtract the time of a baseline loop. The cycle counts of the baseline (6) and of 3 digamma + beta (310) show that the compiler removed no call.

Function level, net ns per call:

kernel	develop	branch	ratio
grad_reg_inc_beta double, incl. 3 digamma + beta	28 022	922	30x
grad_reg_inc_beta var, value + grad	157 825	9 670	16x
grad_reg_inc_beta fvar	88 108	1 776	50x
grad_reg_inc_beta fvar<fvar>	255 280	3 730	68x
inc_beta_dda double, incl. 2 digamma	1 880	866	2.2x

Distribution level, N = 1000 observations, shared var shapes, value plus reverse sweep, ns per observation:

function	shapes	develop	branch	ratio
beta_cdf	5, 5	1 682	1 107	1.5x
beta_lcdf	5, 5	25 822	642	40x
beta_lccdf	5, 5	25 810	642	40x
beta_cdf	50, 50	6 385	2 653	2.4x
beta_lcdf	50, 50	21 295	1 608	13x
beta_lccdf	50, 50	21 297	1 607	13x
neg_binomial_cdf	5, 0.3	1 895	1 724	1.1x
neg_binomial_lcdf	5, 0.3	7 677	983	7.8x
student_t_lcdf	ν = 3	4 938	1 051	4.7x

On develop, path B evaluates the 2F1 power series with three exp and three log calls per term and no reflection. At z = 0.998 it needs 607 607 terms. The continued fraction of the branch costs about 600 cycles in addition to the digamma calls.

Testing

Six test files use fixed references from the 80-digit mpmath computation. None uses finite differences, because finite differences cannot see these defects.

• prim/fun/grad_reg_inc_beta_test.cpp: the two existing cases, with the z = 1 expectation corrected to 0 and the tolerance tightened to 1e-12. New cases for each defect class: beta(a, b) underflow, prefactor underflow, integer b, and the tails.
• prim/fun/inc_beta_dda_test.cpp, prim/fun/inc_beta_ddb_test.cpp: rewritten from the reference at the same twelve points, with the digamma_b argument corrected, plus regression points.
• rev/fun/inc_beta_test.cpp (new): shape gradients of inc_beta with var arguments at the defect points.
• rev/prob/beta_cdf_test.cpp (new): beta_cdf shape partials at the b ≫ a point, with both shapes as autodiff variables and with each shape alone.
• rev/prob/beta_cdf_log_test.cpp (new): beta_lcdf shape gradients at the NaN point, at the dropped-term point and at an integer shape.

Regression check: I installed the develop headers and ran these six files. All 6 suites fail, with 13 of 15 cases failing. With the headers of this branch, all 15 cases pass, both before and after the swap. The two cases that pass on develop are at points where develop was already correct.

Licence note

The new body of grad_reg_inc_beta.hpp derives from TensorFlow Probability, which is Apache-2.0. The file carries the TFP copyright notice and refers to licenses/tensorflow-probability-license.txt. It uses the same form as the Boost-derived code in prim/fun/log_modified_bessel_first_kind.hpp. The tree already contains Apache-2.0 code in vendored libraries (lib/tbb_2020.3, lib/benchmark_1.5.1). This file would be the first Apache-2.0-derived source file under stan/.

Known limits, not addressed here

• beta_lccdf computes 1 - Pn in linear space after inc_beta. At (352, 590, 0.758) the double I is exactly 1, so the value is -inf and the gradients are infinite. The true value is log(1 - I) ≈ -0.34. This defect is in the value code on both trees and needs a separate change (the symmetry relation for the value): Evaluate the incomplete beta complement without cancellation in the lccdf and Student t cdf functions #3433, stacked on this PR.
• Where I is below about 1e-315 (subnormal), the relative error of the gradient rises to 1e-2. The reference itself has few correct digits there.

Not included

• No change to grad_2F1, which other functions use.
• No change to how the neg_binomial*_cdf functions call the root. They need one derivative and already call the root once.

Release notes

Replaces the incomplete beta gradient roots with a more stable continued-fraction implementation

Checklist

• Copyright holder: Aalto University

  The copyright holder is typically you or your assignee, such as a university or company. By submitting this pull request, the copyright holder is agreeing to the license the submitted work under the following licenses:
  - Code: BSD 3-clause (https://opensource.org/licenses/BSD-3-Clause)
  - Documentation: CC-BY 4.0 (https://creativecommons.org/licenses/by/4.0/)

• the basic tests are passing

  • unit tests pass (to run, use: ./runTests.py test/unit)
  • header checks pass, (make test-headers)
  • dependencies checks pass, (make test-math-dependencies)
  • docs build, (make doxygen)
  • code passes the built in C++ standards checks (make cpplint)

• the code is written in idiomatic C++ and changes are documented in the doxygen

• the new changes are tested

[avehtari]

I'll be away for the next two weeks, but made this PR to avoid duplicated work

[avehtari]

The improved grad_reg_inc_beta from this PR is used to fix also other distribution functions in #3433

[avehtari]

@SteveBronder, excellent additional speedup. My measurements

case	original PR/steve
first-order reverse mode, both shapes	1.02–1.10
one shape only (beta_lcdf alpha only, neg_binomial_lcdf)	1.24–1.38
neg_binomial_2_cdf	5.1
fvar<var> cdfs (one Hessian-vector product)	3.0–9.7
fvar<var> inc_beta	1.6

[avehtari]

@SteveBronder

1. The generated test/prob tests found one regression from your changes. beta_proportion_lcdf and beta_proportion_lccdf fail with fvar<fvar<var>> arguments: ibeta_derivative<long double>(...): Overflow Error. The cause is in inc_beta, not in these two files. The new rev and fwd inc_beta compute the z partial with inc_beta_ddz, which for double is Boost's ibeta_derivative. At z = 1 with b < 1 (and at z = 0 with a < 1) the density is infinite, and Boost throws. The previous code used the pow product there, which gave inf without a throw. Three direct reproductions on the head:

• inc_beta(var 1.05, var 0.35, var 1.0), then grad(): throws.
• beta_proportion_lcdf(1.0, fvar<var> mu = 0.75, 1.4), then d_.grad(): throws.
• inc_beta(fvar<double>) at z = 1 with a zero tangent on z: throws.

I pushed a34e6f7. It calls ibeta_derivative with a policy that returns infinity on overflow, and keeps every other default. The result is bit-identical to the default call at 180 360 grid points, and is +inf at the 3 240 points where the default throws. It is the same call, so the speed does not change. At z = 1 the shape gradients of inc_beta are now 0 and the z gradient is +inf. New tests are in prim/fun/inc_beta_ddz_test, rev/fun/inc_beta_test and mix/fun/inc_beta_test; all three fail with the previous inc_beta_ddz.hpp.

2. In the generated tests, the UpperBound case leaves nodes at y = 1 on the stack, and a later grad() in the same process runs them. That is why the same binary passes when a filter skips the earlier cases.

The second commit, 3fd3fe2, starts the new mix test from an empty autodiff stack, because the expect_ad NaN cases earlier in that file leave nodes on the stack whose reverse pass throws.

With these two commits, the generated cdf, lcdf and lccdf tests of beta, beta_proportion, binomial, neg_binomial, neg_binomial_2 and student_t pass (185 programs, --run-all), and so do the unit suites of the incomplete beta functions. Tested on Roihu, AMD EPYC 9965, g++ 11.5.

One behavior stays as before: with a zero tangent on z, the fvar tangent of inc_beta at z = 1 is 0 * inf = NaN.

[andrjohns]

@avehtari can you split this one up as well? Given the amount of code changed, might be better to land this in parts. Thegrad_reg_inc_beta changes are the ones which need the most review/testing, while the distributional and other downstream function changes can land in separate PRs - much easier to review that way

[avehtari]

@andrjohns how many PRs would you prefer? There are 3 logical parts and thus it would be possible to split this in 3 PRs

part	content	base	files	lines
B1 (#3409 reused)	prim grad_reg_inc_beta (continued fraction, GRAD_RIB API), inc_beta_dda, inc_beta_ddb, inc_beta_ddz (with a34e6f7499), fwd inc_beta, prim/fun.hpp, their prim tests	develop	10	+870 -311
B2	rev grad_reg_inc_beta (Hessian-vector product pass, SteveBronder's code), rev inc_beta, rev/fun.hpp, rev and mix tests (with 3fd3fe2f6a)	B1	6	+975 -161
B3	the downstream cdfs on the new API: beta_cdf/lcdf/lccdf, beta_proportion_lcdf/lccdf, neg_binomial_cdf/lcdf/lccdf, neg_binomial_2_cdf, student_t_cdf/lcdf/lccdf, beta cdf tests	B1	14	+369 -528

B2 and B3 are independent of each other. B3 is is big in lines, but it's the same edit in each file.
In addition #3433 follows B3

Instead of 3 PRs, there could be 2 by having B1 and B2 in one PR. Or there could be 5 PRs by splitting B3 by family.

Let me know what you prefer, and I'll do the split

[stan-buildbot]

Name	Old Result	New Result	Ratio	Performance change( 1 - new / old )
stat_comp_benchmarks/benchmarks/gp_regr/gen_gp_data.stan	0.04	0.04	0.98	-1.71% slower
stat_comp_benchmarks/benchmarks/gp_regr/gp_regr.stan	0.17	0.16	1.04	3.52% faster
stat_comp_benchmarks/benchmarks/garch/garch.stan	0.88	0.86	1.02	2.24% faster
stat_comp_benchmarks/benchmarks/low_dim_corr_gauss/low_dim_corr_gauss.stan	0.01	0.01	1.0	-0.21% slower
stat_comp_benchmarks/benchmarks/eight_schools/eight_schools.stan	0.1	0.11	1.0	-0.25% slower
stat_comp_benchmarks/benchmarks/low_dim_gauss_mix/low_dim_gauss_mix.stan	8.23	8.0	1.03	2.77% faster
stat_comp_benchmarks/benchmarks/gp_pois_regr/gp_pois_regr.stan	3.67	3.67	1.0	0.15% faster
stat_comp_benchmarks/benchmarks/pkpd/one_comp_mm_elim_abs.stan	32.97	32.58	1.01	1.16% faster
stat_comp_benchmarks/benchmarks/pkpd/sim_one_comp_mm_elim_abs.stan	0.38	0.38	1.02	1.86% faster
stat_comp_benchmarks/benchmarks/irt_2pl/irt_2pl.stan	6.78	6.86	0.99	-1.17% slower
stat_comp_benchmarks/benchmarks/sir/sir.stan	165.14	166.69	0.99	-0.94% slower
stat_comp_benchmarks/benchmarks/arK/arK.stan	3.23	3.25	0.99	-0.71% slower
stat_comp_benchmarks/benchmarks/arma/arma.stan	0.62	0.62	1.0	0.27% faster
stat_comp_benchmarks/benchmarks/low_dim_gauss_mix_collapse/low_dim_gauss_mix_collapse.stan	25.76	25.72	1.0	0.16% faster
performance.compilation	391.31	386.77	1.01	1.16% faster
Mean result: 1.0057822753525654

Jenkins Console Log
Jenkins Build Stages
Commit hash: 8774185548440f572c1900f4a344e9ba089df121

Machine information

Distributor ID:	Ubuntu
Description:	Ubuntu 20.04.3 LTS
Release:	20.04
Codename:	focal

CPU:

Architecture:                            x86_64
CPU op-mode(s):                          32-bit, 64-bit
Byte Order:                              Little Endian
Address sizes:                           52 bits physical, 57 bits virtual
CPU(s):                                  192
On-line CPU(s) list:                     0-191
Thread(s) per core:                      2
Core(s) per socket:                      48
Socket(s):                               2
NUMA node(s):                            2
Vendor ID:                               AuthenticAMD
CPU family:                              25
Model:                                   17
Model name:                              AMD EPYC 9474F 48-Core Processor
Stepping:                                1
Frequency boost:                         enabled
CPU MHz:                                 1497.134
CPU max MHz:                             4114.4229
CPU min MHz:                             1500.0000
BogoMIPS:                                7189.39
Virtualization:                          AMD-V
L1d cache:                               3 MiB
L1i cache:                               3 MiB
L2 cache:                                96 MiB
L3 cache:                                512 MiB
NUMA node0 CPU(s):                       0-47,96-143
NUMA node1 CPU(s):                       48-95,144-191
Vulnerability Gather data sampling:      Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit:             Not affected
Vulnerability L1tf:                      Not affected
Vulnerability Mds:                       Not affected
Vulnerability Meltdown:                  Not affected
Vulnerability Mmio stale data:           Not affected
Vulnerability Old microcode:             Not affected
Vulnerability Reg file data sampling:    Not affected
Vulnerability Retbleed:                  Not affected
Vulnerability Spec rstack overflow:      Mitigation; Safe RET
Vulnerability Spec store bypass:         Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:                Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:                Mitigation; Enhanced / Automatic IBRS; IBPB conditional; STIBP always-on; PBRSB-eIBRS Not affected; BHI Not affected
Vulnerability Srbds:                     Not affected
Vulnerability Tsa:                       Mitigation; Clear CPU buffers
Vulnerability Tsx async abort:           Not affected
Vulnerability Vmscape:                   Mitigation; IBPB before exit to userspace
Flags:                                   fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good amd_lbr_v2 nopl xtopology nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 hw_pstate ssbd mba perfmon_v2 ibrs ibpb stibp ibrs_enhanced vmmcall fsgsbase bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local user_shstk avx512_bf16 clzero irperf xsaveerptr rdpru wbnoinvd amd_ppin cppc arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif x2avic v_spec_ctrl vnmi avx512vbmi umip pku ospke avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq la57 rdpid overflow_recov succor smca fsrm flush_l1d debug_swap

G++:

g++ (Ubuntu 9.4.0-1ubuntu1~20.04) 9.4.0
Copyright (C) 2019 Free Software Foundation, Inc.
This is free software; see the source for copying conditions.  There is NO
warranty; not even for MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.

Clang:

clang version 10.0.0-4ubuntu1 
Target: x86_64-pc-linux-gnu
Thread model: posix
InstalledDir: /usr/bin

[SteveBronder]

@andrjohns sorry I made this PR a lot longer as I started fixing things adjacent / related to this PR. Some of this code has not been updated in years so I figured while we are here I would fix up nearby places

With my fixes as well a lot of cdfs are now 10x faster than they were before.

I've read over all of Aki's code and can handle the rest of the review if you are okay with that

[stan-buildbot]

Name	Old Result	New Result	Ratio	Performance change( 1 - new / old )
stat_comp_benchmarks/benchmarks/low_dim_gauss_mix_collapse/low_dim_gauss_mix_collapse.stan	8.31	7.71	1.08	7.23% faster
stat_comp_benchmarks/benchmarks/gp_pois_regr/gp_pois_regr.stan	2.13	2.02	1.05	4.85% faster
stat_comp_benchmarks/benchmarks/eight_schools/eight_schools.stan	0.06	0.05	1.06	5.51% faster
stat_comp_benchmarks/benchmarks/irt_2pl/irt_2pl.stan	4.17	3.95	1.06	5.4% faster
stat_comp_benchmarks/benchmarks/garch/garch.stan	0.42	0.4	1.04	4.19% faster
stat_comp_benchmarks/benchmarks/low_dim_gauss_mix/low_dim_gauss_mix.stan	2.65	2.51	1.05	5.14% faster
stat_comp_benchmarks/benchmarks/pkpd/sim_one_comp_mm_elim_abs.stan	0.24	0.23	1.06	5.44% faster
stat_comp_benchmarks/benchmarks/pkpd/one_comp_mm_elim_abs.stan	19.0	17.77	1.07	6.48% faster
stat_comp_benchmarks/benchmarks/gp_regr/gen_gp_data.stan	0.02	0.02	1.04	3.4% faster
stat_comp_benchmarks/benchmarks/gp_regr/gp_regr.stan	0.1	0.09	1.08	6.99% faster
stat_comp_benchmarks/benchmarks/sir/sir.stan	88.09	82.25	1.07	6.62% faster
stat_comp_benchmarks/benchmarks/arK/arK.stan	1.54	1.45	1.06	5.88% faster
stat_comp_benchmarks/benchmarks/low_dim_corr_gauss/low_dim_corr_gauss.stan	0.01	0.01	1.02	2.29% faster
stat_comp_benchmarks/benchmarks/arma/arma.stan	0.28	0.26	1.07	6.39% faster
performance.compilation	226.78	226.84	1.0	-0.03% slower
Mean result: 1.053612264726441

Jenkins Console Log
Jenkins Build Stages
Commit hash: f650e7d0b68a35569be5874169d8951cbf1b6ea6

Machine information

Distributor ID:	Ubuntu
Description:	Ubuntu 20.04.3 LTS
Release:	20.04
Codename:	focal

CPU:

Architecture:                            x86_64
CPU op-mode(s):                          32-bit, 64-bit
Byte Order:                              Little Endian
Address sizes:                           46 bits physical, 48 bits virtual
CPU(s):                                  80
On-line CPU(s) list:                     0-79
Thread(s) per core:                      2
Core(s) per socket:                      20
Socket(s):                               2
NUMA node(s):                            2
Vendor ID:                               GenuineIntel
CPU family:                              6
Model:                                   85
Model name:                              Intel(R) Xeon(R) Gold 6248 CPU @ 2.50GHz
Stepping:                                7
CPU MHz:                                 1001.187
CPU max MHz:                             3900.0000
CPU min MHz:                             1000.0000
BogoMIPS:                                5000.00
Virtualization:                          VT-x
L1d cache:                               1.3 MiB
L1i cache:                               1.3 MiB
L2 cache:                                40 MiB
L3 cache:                                55 MiB
NUMA node0 CPU(s):                       0,2,4,6,8,10,12,14,16,18,20,22,24,26,28,30,32,34,36,38,40,42,44,46,48,50,52,54,56,58,60,62,64,66,68,70,72,74,76,78
NUMA node1 CPU(s):                       1,3,5,7,9,11,13,15,17,19,21,23,25,27,29,31,33,35,37,39,41,43,45,47,49,51,53,55,57,59,61,63,65,67,69,71,73,75,77,79
Vulnerability Gather data sampling:      Mitigation; Microcode
Vulnerability Indirect target selection: Mitigation; Aligned branch/return thunks
Vulnerability Itlb multihit:             KVM: Mitigation: Split huge pages
Vulnerability L1tf:                      Not affected
Vulnerability Mds:                       Not affected
Vulnerability Meltdown:                  Not affected
Vulnerability Mmio stale data:           Mitigation; Clear CPU buffers; SMT vulnerable
Vulnerability Old microcode:             Not affected
Vulnerability Reg file data sampling:    Not affected
Vulnerability Retbleed:                  Mitigation; Enhanced IBRS
Vulnerability Spec rstack overflow:      Not affected
Vulnerability Spec store bypass:         Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:                Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:                Mitigation; Enhanced / Automatic IBRS; IBPB conditional; PBRSB-eIBRS SW sequence; BHI SW loop, KVM SW loop
Vulnerability Srbds:                     Not affected
Vulnerability Tsa:                       Not affected
Vulnerability Tsx async abort:           Mitigation; TSX disabled
Vulnerability Vmscape:                   Mitigation; IBPB before exit to userspace
Flags:                                   fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf pni pclmulqdq dtes64 monitor ds_cpl vmx smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid dca sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault epb cat_l3 cdp_l3 intel_ppin ssbd mba ibrs ibpb stibp ibrs_enhanced tpr_shadow flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid cqm mpx rdt_a avx512f avx512dq rdseed adx smap clflushopt clwb intel_pt avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local dtherm ida arat pln pts hwp hwp_act_window hwp_epp hwp_pkg_req vnmi pku ospke avx512_vnni md_clear flush_l1d arch_capabilities

G++:

g++ (Ubuntu 9.4.0-1ubuntu1~20.04) 9.4.0
Copyright (C) 2019 Free Software Foundation, Inc.
This is free software; see the source for copying conditions.  There is NO
warranty; not even for MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.

Clang:

clang version 10.0.0-4ubuntu1 
Target: x86_64-pc-linux-gnu
Thread model: posix
InstalledDir: /usr/bin

[andrjohns]

No worries both! I think it would be good to split at least the B1 row from Aki's table to a separate PR, since those are the big algorithmic changes