Skip to content

feat(mlflow): autolog support for BassModel#2624

Open
anevolbap wants to merge 3 commits into
pymc-labs:mainfrom
anevolbap:feat/2588-mlflow-bass
Open

feat(mlflow): autolog support for BassModel#2624
anevolbap wants to merge 3 commits into
pymc-labs:mainfrom
anevolbap:feat/2588-mlflow-bass

Conversation

@anevolbap

@anevolbap anevolbap commented Jun 12, 2026

Copy link
Copy Markdown
Contributor

Closes #2588

Adds Bass model support to pymc_marketing.mlflow.autolog, following the existing MMM and CLV patterns:

  • New log_bass flag (default True) patches BassModel.fit to log the model configuration, stamp the run id on idata.attrs, and store the idata.nc artifact.
  • New log_bass_configuration helper logs model type, version, sampler config, and the prior configuration (m, p, q, likelihood).
  • Module and autolog docstrings document the new integration with an example.

While testing this, autolog on a default-prior BassModel failed about half the time with ValueError: lam < 0 or lam contains NaNs. The cause is pm.model_to_graphviz, which draws from the prior graph to evaluate plate shapes; the default m ~ Normal(0, 10) prior then produces a negative Poisson rate. The same happens with pm.sample_prior_predictive on a default BassModel, so the root cause is the prior itself: should the default prior for m be positive (HalfNormal, LogNormal, truncated Normal)? Happy to open a separate issue for that. Meanwhile, a separate commit makes log_model_graph treat 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.sample that BassModel.fit calls internally; this PR only adds the Bass-specific layer described in the issue.

@github-actions github-actions Bot added tests mlflow Bass model Dealing with the Bass Defusion model docs Improvements or additions to documentation enhancement New feature or request labels Jun 12, 2026
@codecov

codecov Bot commented Jun 12, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 94.20%. Comparing base (649547e) to head (12e0851).

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

@anevolbap anevolbap force-pushed the feat/2588-mlflow-bass branch from 15f91d4 to 12e0851 Compare June 12, 2026 20:40
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Bass model Dealing with the Bass Defusion model docs Improvements or additions to documentation enhancement New feature or request mlflow tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

MLflow integration for BassModel

1 participant