Fix gradient-aggregation bias in GradientMode parameter updates

[giovanni-soma]

Summary

GradientMode (and GradientMPIMode) computed a biased parameter gradient for any batch with more than one distinct sample.

The optimization layer (JuMP + DiffOpt) returns dCdy, the gradient of the assessment cost w.r.t. the forecasts. To update the network parameters we propagate these through the forecast model with a linear surrogate loss whose θ-gradient equals the chain-rule term. The bug: dCdy was aggregated across samples before building the surrogate, so the update computed

mean_t(dCdy_t) · mean_t(∂ŷ_t/∂θ)

instead of the correct empirical-cost gradient

(1/T) Σ_t dCdy_t · ∂ŷ_t/∂θ

The two differ by the sample covariance between the per-sample cost gradient and the prediction Jacobian — zero only when all samples are identical or T = 1. Every existing test used T = 1, so the bug was latent.

Changes

• compute_cost: no longer aggregates dC. With with_gradients=true it
  returns the full per-sample gradient matrix of size (T, output_size). The
  aggregate flag now affects only the cost.
• apply_gradient!: signature changed from dCdy::AbstractVector to
  dCdy::AbstractMatrix (T, output_size); the surrogate loss now contracts
  each sample's gradient with its own forecast:
  loss = sum(dCdy' .* m(X')) / T.
• gradient_mpi.jl: the MPI path had the same aggregation; per-sample
  gradients are now stacked into a row-aligned matrix, and the g_tol stop rule
  / apply_gradient! call updated accordingly.
• Tests: added GradientMode Multi-Sample Gradient Correctness, which checks
  the gradient the framework applies against a central finite-difference
  reference of the true cost (it fails on the old behavior: applied -2.0 vs
  true -1.0). Updated test_predictive_model.jl calls to the new matrix API.
• Docs: clarified the g_tol description in modes.md.
• Version bump 0.1.6 → 0.1.7.

Compatibility

⚠️ Technically breaking: compute_cost(...; with_gradients=true) now returns a matrix instead of an aggregated vector, and apply_gradient! takes a matrix. Released as a patch (0.1.7) given the limited public footprint.

Testing

Full suite passes locally (Julia 1.11), including the new regression test. JuliaFormatter (v1) clean.

[codecov]

Codecov Report

❌ Patch coverage is 20.00000% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 67.40%. Comparing base (a324bd7) to head (242e05f).
⚠️ Report is 20 commits behind head on dev.

Files with missing lines	Patch %	Lines
src/optimizers/gradient_mpi.jl	0.00%	4 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##              dev      #56      +/-   ##
==========================================
- Coverage   67.45%   67.40%   -0.06%     
==========================================
  Files          12       12              
  Lines         590      589       -1     
==========================================
- Hits          398      397       -1     
  Misses        192      192              

☔ View full report in Codecov by Sentry.
📢 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.