Vectorize log score gradient calculation

[VisruthSK]

This PR was made with assistance from Codex.

Swaps a gradient calculation in stacking weights from a loop to matrix ops. Adds a test for correctness before changing implementation.

Benchmark

set.seed(0)
N <- 400
K <- 20
lpd_point <- matrix(rnorm(N * K, mean = -1), nrow = N, ncol = K)
w <- rep(1 / K, K - 1)

stacking_gradient_loop <- function(w, lpd_point) {
  K <- ncol(lpd_point)
  w_full <- c(w, 1 - sum(w))
  grad <- rep(0, K - 1)
  mlpd <- matrixStats::rowMaxs(lpd_point)

  for (k in 1:(K - 1)) {
    grad[k] <- sum(
      (exp(lpd_point[, k] - mlpd) -
        exp(lpd_point[, K] - mlpd)) /
        exp(
          matrixStats::rowLogSumExps(sweep(lpd_point, 2, log(w_full), '+')) -
            mlpd
        )
    )
  }

  -grad
}

stacking_gradient_vectorized <- function(w, lpd_point) {
  K <- ncol(lpd_point)
  w_full <- c(w, 1 - sum(w))
  mlpd <- matrixStats::rowMaxs(lpd_point)
  denom <- exp(
    matrixStats::rowLogSumExps(sweep(lpd_point, 2, log(w_full), '+')) - mlpd
  )
  grad <- colSums(
    (exp(lpd_point[, 1:(K - 1), drop = FALSE] - mlpd) -
      exp(lpd_point[, K] - mlpd)) /
      denom
  )

  -grad
}

results <- bench::mark(
  loop = stacking_gradient_loop(w, lpd_point),
  vectorized = stacking_gradient_vectorized(w, lpd_point),
  iterations = 1e3
)

summary(results)
#> # A tibble: 2 × 6
#>   expression      min   median `itr/sec` mem_alloc `gc/sec`
#>   <bch:expr> <bch:tm> <bch:tm>     <dbl> <bch:byt>    <dbl>
#> 1 loop         6.48ms    7.6ms      125.    7.26MB     7.15
#> 2 vectorized  624.4µs  872.9µs      810.  283.69KB     3.25
summary(results, relative = TRUE)
#> # A tibble: 2 × 6
#>   expression   min median `itr/sec` mem_alloc `gc/sec`
#>   <bch:expr> <dbl>  <dbl>     <dbl>     <dbl>    <dbl>
#> 1 loop        10.4   8.71      1         26.2     2.20
#> 2 vectorized   1     1         6.46       1       1

^{Created on 2026-06-06 with reprex v2.1.1}

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 92.70%. Comparing base (286263d) to head (2a09e6d).
⚠️ Report is 2 commits behind head on master.

Additional details and impacted files

@@           Coverage Diff           @@
##           master     #370   +/-   ##
=======================================
  Coverage   92.70%   92.70%           
=======================================
  Files          31       31           
  Lines        3029     3031    +2     
=======================================
+ Hits         2808     2810    +2     
  Misses        221      221           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

This is how benchmark results would change (along with a 95% confidence interval in relative change) if 2a09e6d is merged into master:

• ✔️loo_function: 1.76s -> 1.77s [-0.1%, +0.52%]
• ✔️loo_matrix: 1.33s -> 1.34s [-0.17%, +0.98%]
  Further explanation regarding interpretation and methodology can be found in the documentation.