Cpaniaguam/fix rl likelihood builder pymc6

[cpaniaguam]

Prevent PyTensor from incorrectly optimizing JAX-backed operations.

JAX-PyTensor Integration Improvements:

• Added do_constant_folding methods to both LANLogpOp and LANLogpVJPOp classes in src/hssm/distribution_utils/jax.py to prevent PyTensor from attempting to precompute (constant fold) outputs of JAX-backed operations, ensuring correct runtime behavior. [1] [2]

Test Suite Updates:

• Updated test_make_rl_logp_op in tests/rl/test_rl_likelihood_builder.py to use pytensor.function with mode="FAST_COMPILE" for gradient evaluation, improving test reliability and compatibility.
• Marked test_predictive_idata_to_dataframe in tests/test_utils.py with @pytest.mark.xfail to indicate it is expected to fail due to recent changes in PyMC, preventing it from causing false negatives in CI. See Cpaniaguam/fix predictive idata to dataframe datatree #974 with a fix.

[Copilot]

Pull request overview

This PR improves the PyTensor↔JAX integration by preventing PyTensor from constant-folding JAX-backed Ops (which can lead to incorrect/unsupported compile-time evaluation), and adjusts tests to be more robust under recent PyTensor/PyMC behavior changes.

Changes:

• Disabled PyTensor constant folding for the JAX-backed LANLogpOp and LANLogpVJPOp.
• Updated the RL likelihood builder gradient test to evaluate gradients via a compiled pytensor.function(..., mode="FAST_COMPILE").
• Marked test_predictive_idata_to_dataframe as xfail due to upstream PyMC changes.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

File	Description
src/hssm/distribution_utils/jax.py	Prevents constant folding of JAX-backed Ops to avoid incorrect compile-time evaluation.
tests/rl/test_rl_likelihood_builder.py	Uses a compiled PyTensor function in FAST_COMPILE mode for gradient evaluation reliability.
tests/test_utils.py	Marks a known-broken test as xfail to avoid CI false negatives.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

[digicosmos86]

Why do we need these packages? Are these absolutely necessary to maintain basic functionalities?

[cpaniaguam]

I think so. HSSM still restores .nc traces through az.from_netcdf(...)

HSSM/src/hssm/base.py

Line 1490 in 95b27df

traces = az.from_netcdf(traces)

HSSM/src/hssm/base.py

Line 1508 in 95b27df

traces = az.from_netcdf(traces)

HSSM/src/hssm/base.py

Line 1665 in 95b27df

idata_dict["idata_mcmc"] = az.from_netcdf(traces_path)

What seems to have changed is that newer xarray/ArviZ installs no longer seem to guarantee a working NetCDF backend transitively, and CI started failing when pytest tried to load the .nc fixture in

HSSM/tests/conftest.py

Line 91 in 95b27df

return az.from_netcdf("tests/fixtures/cavanagh_idata.nc")

[digicosmos86]

Yeah it seems that both xarray and arviz now have them as optional and/or dev dependencies. I think these packages do that for a good reason, and we probably should do the same and then add these packages to dev dependencies. In code, we can probably check if these packages are installed and throw an informative error if not.

I like to keep the dependencies as slim as possible and not add additional packages for things that are occasionally used. @cpaniaguam @AlexanderFengler @krishnbera thoughts?

@cpaniaguam these .nc files might need to be generated again since they came from old InferenceData saves

[digicosmos86]

Additionally, I think it could be a good idea to have an io optional group and have cloudpickle, h5py and h5netcdf in that group

[digicosmos86]

@cpaniaguam the upstream errors should all be fixed now. Can you change the base to 996-fix-model-cartoon-plot and see if all tests pass now?

[cpaniaguam]

@cpaniaguam the upstream errors should all be fixed now. Can you change the base to 996-fix-model-cartoon-plot and see if all tests pass now?

Sure!

[cpaniaguam]

At least one test is failing https://github.com/lnccbrown/HSSM/actions/runs/27966221081/job/82760632068
@digicosmos86

[coderabbitai]

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 4c8f4750-75cd-4c04-bffa-b7c85f6bb1d9

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch cpaniaguam/fix-rl-likelihood-builder-pymc6

---

_{Comment @coderabbitai help to get the list of available commands.}

[cpaniaguam]

@digicosmos86 There still is one function that needs to be updated to handle xr.DataTree in src/hssm/utils.py -- predictive_idata_to_dataframe. The corresponding test is in tests/test_utils.py::test_predictive_idata_to_dataframe. I kept the failing marking on it.

[digicosmos86]

@digicosmos86 There still is one function that needs to be updated to handle xr.DataTree in src/hssm/utils.py -- predictive_idata_to_dataframe. The corresponding test is in tests/test_utils.py::test_predictive_idata_to_dataframe. I kept the failing marking on it.

This is fixed here https://github.com/lnccbrown/HSSM/pull/993/changes#diff-11c83362802db3028664e029ce20f83d181667e3631fe618614ad9672665d383R558