Add retry logic for initial parameters in ExternalSampler (#2739)

[0AnshuAditya0]

Fixes #2739

This PR adds retry logic when generating initial parameters for external samplers, similar to what HMC already does.

Changes:

• Modified src/mcmc/external_sampler.jl to attempt finding valid initial parameters up to 10 times
• Checks that both logp and gradient are finite before accepting initial parameters
• Provides informative error message if all attempts fail

Why:
External samplers previously only tried once to generate initial parameters. If the first attempt produced infinite logp or gradient, sampling would fail immediately. This makes external samplers more robust.

[mhauru]

Hi @0AnshuAditya0, thanks for your interest in contributing to Turing.jl.

Some of our team is currently on leave, and I'm really busy and will soon go on leave, so it may be early January before we'll have time to give your PR a proper look. So apologies if we aren't the most responsive right now, we'll get to it in a few weeks.

Some quick first thoughts though:

• Rather than writing two parallel versions of this, for HMC specifically in hmc.jl and for external samplers in external_samplers.jl, I think we should try to unify these under a single function. Would you be happy to refactor so that there would be a single function called find_initial_params that handles both cases? On a quick look, I think the only part of the current find_initial_params function that is specific to HMC is the check of whether the phasepoint is finite, but that's just checking the finiteness of the value and the gradient, so that function could easily be modified to handle the general case.
• Not all samplers require gradients, and not all models have well-defined gradients, so we should have a flag for whether find_initial_params checks for gradient finiteness at all. I think by default this should be false.
• We should have a test for this. I'm not sure how to best do that. Maybe a model which is manually crafted to give infinite log probabilities or infinite gradients e.g. 90% of the time? And then testing on that to see that we can still eventually get a good initial value. The test model could be something very artificial, doesn't have to be a meaningful statistical model at all.

[0AnshuAditya0]

Thanks for the feedback @mhauru! I appreciate the guidance on refactoring this properly.

Your approach makes a lot of sense. Instead of duplicating the retry logic, I can extract it into a shared utility that handles the general case. Here's what I'm thinking:

Proposed changes:

1. Create a new find_initial_params function in src/mcmc/initial_params.jl that centralizes the retry logic for generating valid initial parameters
2. Add a check_gradient::Bool flag (defaulting to false) so gradient-free samplers can use it without unnecessary gradient evaluations
3. Refactor both hmc.jl and external_sampler.jl to use this shared function, which should eliminate code duplication and make future maintenance easier
4. Add a test using a deliberately problematic model (perhaps one with a constrained parameter space that makes most random initializations invalid) to verify the retry mechanism works reliably

The shared function would handle:

• Multiple initialization attempts (configurable max_attempts)
• Checking finiteness of logp (always)
• Optionally checking gradient finiteness (when check_gradient=true)
• Parameter linking/unlinking for constrained spaces
• Clear error messages when all attempts are exhausted

Does this implementation plan align with what you had in mind? Happy to adjust the approach if you have other preferences.

I'll work on these changes over the next few days. No rush on the review - I understand the holiday season timing. Take your time, and I'll have the refactored version ready when you're back!

[0AnshuAditya0]

Hi @mhauru! I've implemented the changes you suggested. Here's what I did:

Changes Made

1. Created Shared Function

• Added src/mcmc/initial_params.jl with find_initial_params() function
• Both HMC and external_sampler now use this shared function
• Used a validator pattern where each sampler passes its own validation logic

2. Made Gradient Checking Configurable

• External samplers check gradients using logdensity_and_gradient()
• HMC uses AHMC.phasepoint() for validation
• Each sampler can define what "valid parameters" means

3. Added Tests

• Created test/mcmc/initial_params.jl with comprehensive tests
• Tests include models that fail initialization ~90% of the time
• Verified retry logic works for both NUTS and HMC
• Tests pass locally ✓

Why Validator Pattern?

I chose this approach because HMC and external samplers validate parameters differently:

• HMC: validates using AHMC.phasepoint()
• External: validates using LogDensityProblems interface

The validator pattern keeps them flexible while sharing the retry logic. Let me know if a simpler approach would be better!

Testing

Tested with constrained models locally - the retry logic works as expected.

Please let me know if you'd like me to change anything. Thanks for the guidance on this!

[penelopeysm]

Thank you for the PR @0AnshuAditya0, I have a number of comments above 🙂

[0AnshuAditya0]

All feedback addressed! Here's what I changed:

✅ Moved find_initial_params to abstractmcmc.jl (more appropriate location)
✅ Simplified validator to return Bool only (removed diagnostics tuple)
✅ Updated all validator implementations in HMC and external_sampler
✅ Moved tests to test/mcmc/abstractmcmc.jl
✅ Improved test with counter-based model (deterministic 30 failures)
✅ Added direct find_initial_params unit tests
✅ Removed duplicate warning test from test/mcmc/hmc.jl
✅ Removed unnecessary StableRNG usage
✅ Fixed docstring (max_attempts=1000)

Ready for re-review! Thanks for the detailed feedback.

[penelopeysm]

This is looking a lot nicer. More comments below.

Could I also point out, that if I wanted to talk to Claude and ask it to do this PR, I would just ask Claude to do it? I think you are smart enough to use your own words and your own programming skills to work on this. Bear in mind that people are spending real time to review your PRs.

[0AnshuAditya0]

thanks for the feedback.

The change removes a random-value assertion and fixes the RNG seed so the test result is stable and not dependent on sampler randomness. The intent was to keep the behavior the same while making the test deterministic.

I’m still learning and will keep improving how I communicate changes. Thanks for taking the time to review this.

[penelopeysm]

You incorporated a few of my changes in fafdf93 (great!), but then you reverted them in the next commit b676d6a and didn't put them back in again. I think the first commit seems quite reasonable. Is there a reason why you reverted?

[0AnshuAditya0]

i reverted the commit while trying to simplify the tests
are you referring to @addlogprob! (abs(x) < 0.3) ? 0.0 : -Inf
I added probabilistic approach. Should it be deterministic?

[penelopeysm]

The test in the earlier commit was fine before you reverted it. No don't use abs(x) please, the Ref counter was exactly what's needed.

[codecov]

Codecov Report

❌ Patch coverage is 77.77778% with 6 lines in your changes missing coverage. Please review.
✅ Project coverage is 55.74%. Comparing base (4dc7ad0) to head (532aa82).
⚠️ Report is 4 commits behind head on main.

Files with missing lines	Patch %	Lines
src/mcmc/abstractmcmc.jl	53.84%	6 Missing ⚠️

❗ There is a different number of reports uploaded between BASE (4dc7ad0) and HEAD (532aa82). Click for more details.

HEAD has 6 uploads less than BASE

Flag	BASE (4dc7ad0)	HEAD (532aa82)
	24	18

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #2740       +/-   ##
===========================================
- Coverage   86.51%   55.74%   -30.78%     
===========================================
  Files          20       20               
  Lines        1261     1270        +9     
===========================================
- Hits         1091      708      -383     
- Misses        170      562      +392     

☔ 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.

[penelopeysm]

There are still other things from my previous review which were in the reverted commit.