MLflow callback logging

[wd60622]

Description

Add a callback that works with the "pymc" nuts sampler

Related Issue

• Closes #
• Related to Make use of mlflow.log_metric(..., step=..., timestamp=...) #1397

Checklist

• Checked that the pre-commit linting/style checks pass. Feel free to comment pre-commit.ci autofix to auto-fix.
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks) using numpydoc format.
• If you are a pro: each commit corresponds to a relevant logical change

---

📚 Documentation preview 📚: https://pymc-marketing--1425.org.readthedocs.build/en/1425/

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 94.62%. Comparing base (2ef7b51) to head (153bf6f).
Report is 1 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1425      +/-   ##
==========================================
+ Coverage   94.59%   94.62%   +0.02%     
==========================================
  Files          50       50              
  Lines        5388     5416      +28     
==========================================
+ Hits         5097     5125      +28     
  Misses        291      291              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[juanitorduz]

I love this! 🚀 . How ward is to do this for other samplers like nutpie or numpyro?

[wd60622]

I love this! 🚀 . How ward is to do this for other samplers like nutpie or numpyro?

Both need a bit of work. However, there are some conversation happening!

Numpyro status: https://forum.pyro.ai/t/is-callback-supported-for-numpyro-nuts-sampler/6308?u=alfredo_delgado
Nutpie conversation starting here: pymc-devs/nutpie#166