Fix rldm vjp function and update rlssm tutorial

[krishnbera]

This PR fixed the issue in rldm vjp_logp() that caused VI calls (both via HSSM and PyMC) to break.
The rlssm demo tutorial is also updates to showcase a working example with NUTS and VI.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[Copilot]

Pull Request Overview

This PR corrects the vector-Jacobian product function for the RLDM likelihood and updates its documentation and output slicing.

• Changed vjp_logp signature to variadic inputs and keyword-only gz
• Refined the returned gradient slice to exclude data and extra fields
• Updated docstring to describe new inputs format

Comments suppressed due to low confidence (3)

src/hssm/likelihoods/rldm.py:221

• The docstring still refers to inputs as a single list, but the function now uses *inputs varargs. Update the parameter description to match the variadic signature.

        inputs : list

src/hssm/likelihoods/rldm.py:213

• [nitpick] It would be helpful to add unit tests for the updated vjp_logp behavior, especially to verify the new argument handling and the correct slicing of the returned gradients.

        A function that computes the VJP of the log likelihood for the RLDM model.

src/hssm/likelihoods/rldm.py:216

• Changing to *inputs, gz makes gz keyword-only and may break callers passing gz positionally. Consider allowing gz as a positional argument or update all call sites and doc to reflect the new signature.

    def vjp_logp(*inputs, gz):

[Copilot]

Using a hardcoded slice 1:7 can become brittle if the number of parameters or extra fields changes. Consider computing the slice dynamically or documenting why exactly six elements are selected.

Suggested change

	return vjp_logp(gz)[1:7] # Exclude the data and the extra fields
	n_params = 6 # Number of model parameters: rl_alpha, scaler, a, z, t, theta
	return vjp_logp(gz)[1 : 1 + n_params] # Exclude the first field and retain only the parameters

[digicosmos86]

LGTM! I am happy to make n_params an optional parameter to make_vjp_logp so the function is more generalizable

[codecov]

Codecov Report

Attention: Patch coverage is 50.00000% with 1 line in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/hssm/likelihoods/rldm.py	50.00%	1 Missing ⚠️

Files with missing lines	Coverage Δ
src/hssm/likelihoods/rldm.py	96.96% <50.00%> (ø)

🚀 New features to boost your workflow:

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