fix(maths): eliminate cancellation error in exact parameter info-gain weight

[gauraaansh]

Fixes #340

_exact_wnorm computes eq. (D.15) from Da Costa et al. 2020 as:

w = log(sumA) - log(A) + 1/A - 1/sumA + digamma(A) - digamma(sumA)

As A shrinks toward MINVAL, 1/A and digamma(A) both diverge with opposite
sign and are supposed to cancel down to an O(1) value. In float32 that
cancellation is inexact, so the term comes out wrong by several orders of
magnitude for small Dirichlet counts. The error is also erratic rather than
a smooth degradation, matching the oscillating behavior described in the
issue.

Using the digamma recurrence digamma(x+1) = digamma(x) + 1/x, substitute
for 1/A and 1/sumA:

w = log(sumA) - log(A)
    + [digamma(A+1) - digamma(A)]
    - [digamma(sumA+1) - digamma(sumA)]
    + digamma(A) - digamma(sumA)

The digamma(A) and digamma(sumA) terms cancel exactly, leaving:

w = log(sumA) - log(A) + digamma(A+1) - digamma(sumA+1)

This reformulation is algebraically identical to the original expression. It
changes only the numerical evaluation order. The divergent terms are removed
before any floating-point evaluation happens, so the implementation no
longer depends on cancellation between divergent quantities.

Testing

• Randomized stress testing against float64 and mpmath, plus a deterministic
  logarithmic sweep from 1e-20 to 1e3: the original formula's error grows to
  roughly six orders of magnitude near the clipping threshold, while the
  rewritten version stays within float32's ordinary rounding error across the
  same range.
• Verified under jit, vmap, and autodiff (the composition control.py and
  agent.py use for policy inference), including the zero-entry case from
  Bayesian model reduction.
• The existing correctness test previously compared the implementation to a
  hand-copied version of the same pre-fix formula, so it would pass
  regardless of correctness. Added a test against -log(K), the closed-form
  limit eq. (D.15) converges to for a symmetric Dirichlet as concentration
  goes to zero. This reference is independent of either formula.
• Full suite passes (294 passed, 1 skipped, one pre-existing unrelated
  failure from a missing numpyro dependency). ruff check is clean.

[conorheins]

weirdly it appears the 3.14 unit tests are timing out on the infer_states optimization notebooks, and the other 3 tested versions (3.11-3.13) are just being cancelled. will look into this when i have time