feat(mlflow): autolog support for BassModel#2624
Open
anevolbap wants to merge 3 commits into
Open
Conversation
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2624 +/- ##
=======================================
Coverage 94.19% 94.20%
=======================================
Files 100 100
Lines 15384 15403 +19
=======================================
+ Hits 14491 14510 +19
Misses 893 893 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
5 tasks
This was referenced Jun 12, 2026
15f91d4 to
12e0851
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #2588
Adds Bass model support to
pymc_marketing.mlflow.autolog, following the existing MMM and CLV patterns:log_bassflag (defaultTrue) patchesBassModel.fitto log the model configuration, stamp the run id onidata.attrs, and store theidata.ncartifact.log_bass_configurationhelper logs model type, version, sampler config, and the prior configuration (m,p,q,likelihood).autologdocstrings document the new integration with an example.While testing this, autolog on a default-prior
BassModelfailed about half the time withValueError: lam < 0 or lam contains NaNs. The cause ispm.model_to_graphviz, which draws from the prior graph to evaluate plate shapes; the defaultm ~ Normal(0, 10)prior then produces a negative Poisson rate. The same happens withpm.sample_prior_predictiveon a defaultBassModel, so the root cause is the prior itself: should the default prior formbe positive (HalfNormal,LogNormal, truncated Normal)? Happy to open a separate issue for that. Meanwhile, a separate commit makeslog_model_graphtreat graph creation as best-effort, logging a message and continuing, the same way it already handles a missing graphviz install or a render failure.Sampler diagnostics, ArviZ summary, model graph, and metadata logging already worked through the patched
pm.samplethatBassModel.fitcalls internally; this PR only adds the Bass-specific layer described in the issue.