fix: ensure we properly indicate to jax that beta does not have grads for moffat and spergel

[beckermr]

I updated the doc string and I added a stop_gradient call to help JAX along instead of asking it to propagate zeros everywhere.

closes #184

[codspeed-hq]

Merging this PR will not alter performance

⚠️ Unknown Walltime execution environment detected

Using the Walltime instrument on standard Hosted Runners will lead to inconsistent data.

For the most accurate results, we recommend using CodSpeed Macro Runners: bare-metal machines fine-tuned for performance measurement consistency.

✅ 36 untouched benchmarks

---

_{Comparing clarify-moffat-beta-deriv (aa9d01a) with main (51df286)}

[beckermr]

@EiffL can you look at these stop_gradient calls and tell me if they make sense?

[EiffL]

Look reasonable, and I think this should solve the gradient problems, but I'd have to run my test script to be sure, which I can only do tonight (I've been surprised by similar things in the past, like the stop gradient not actually preventing the computation of some gradients on the path).

[beckermr]

It doesn't appear to help the benchmarks so I am guessing the gradient problems are not helped.

[EiffL]

GPU Benchmark: clarify-moffat-beta-deriv vs main vs besssel_improv
┌────────────────┬─────────┬────────────────┬───────────────────────────┐
│ Metric │ main │ besssel_improv │ clarify-moffat-beta-deriv │
├────────────────┼─────────┼────────────────┼───────────────────────────┤
│ Fwd HLO │ 5,472 │ 5,169 │ 5,476 │
├────────────────┼─────────┼────────────────┼───────────────────────────┤
│ Grad HLO │ 6,349 │ 9,384 │ 6,353 │
├────────────────┼─────────┼────────────────┼───────────────────────────┤
│ Fwd time │ 0.647ms │ 1.161ms │ 0.648ms │
├────────────────┼─────────┼────────────────┼───────────────────────────┤
│ Grad time │ 23.0ms │ 135.9ms │ 24.1ms │
├────────────────┼─────────┼────────────────┼───────────────────────────┤
│ Grad/Fwd ratio │ 35.6x │ 117.1x │ 37.3x │
└────────────────┴─────────┴────────────────┴───────────────────────────┘
This branch is essentially identical to main in performance (within noise). The stop_gradient(beta) and
documentation changes don't affect HLO or timing. The besssel_improv regression is not present here since
this branch is based on main's bessel.py.

[EiffL]

doesnt seem to help

[beckermr]

Yep. The benchmark tests on the cpu seem to be a reasonably accurate predictor. We should still merge this pr after a few changes.