Add the ability to compile models to c function pointers with numba

[WardBrian]

Description

This PR upstreams the numba half of nutpie's compile_pymc.py

The code is directly taken from that file, with the only edits being the removal of the JAX portions and a small bugfix

Testing:

Nutpie doesn't provide any direct testing of this functionality besides through calling into the Rust sampler, so I am open for suggestions on how to test this in pymc's own repo.

For what it is worth, my purpose of upstreaming this (use in walnutpie) does work:

import pymc as pm
import pymc.sampling.numba
import walnutpie
import numpy as np


def walnutpie_pymc(model, **kwargs):

    if "num_params" not in kwargs and "inits" not in kwargs:
        chains = kwargs.get("num_chains", 4)
        seed = kwargs.get("seed") or walnutpie.util.rand_u32()
        kwargs['inits'] = np.array([model.initial_point_func(seed + i) for i in range(chains)])

    # TODO: extraction function support in walnutpie
    return walnutpie.walnuts_pyfunc(
        (model.compiled_logp_func.ctypes, model.user_data.ctypes),
        **kwargs,
    )


with pm.Model() as model:
    mu = pm.Data("mu", -0.1)
    sigma = pm.Data("sigma", np.ones(3))
    pm.Normal("a", mu=mu, sigma=sigma, shape=3)

compiled = pymc.sampling.numba.compile_pymc_model(model)
trace = walnutpie_pymc(compiled)
print(walnutpie.mean(trace), walnutpie.standard_deviation(trace), walnutpie.ess(trace))

compiled2 = compiled.with_data(mu=10.0, sigma=3 * np.ones(3))
trace2 = walnutpie_pymc(compiled2)
print(walnutpie.mean(trace2), walnutpie.standard_deviation(trace2), walnutpie.ess(trace2))

Related Issue

• Closes Stable PyMC-to-numba API? nutpie#336
• Related to #

Checklist

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

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

[welcome]

]
💖 Thanks for opening this pull request! 💖 The PyMC community really appreciates your time and effort to contribute to the project. Please make sure you have read our Contributing Guidelines and filled in our pull request template to the best of your ability.

[read-the-docs-community]

Documentation build overview

📚 pymc | 🛠️ Build #34227575 | 📁 Comparing b151d8e against latest (da8fc47)

  🔍 Preview build  

1 file changed

± glossary.html