Cache numba wrappers

[ricardoV94]

This was brought up in #334 but there was a lot more going on there. We should still try to cache identical functions if we can. Code looks over-engineered to me so still a draft

from time import perf_counter

import numpy as np
import pymc as pm
import pytensor

pytensor.config.numba__cache = True

with pm.Model() as model:
    mu = pm.Normal("mu", shape=2)
    y = pm.Data("y", np.arange(100, dtype=float) / 100)
    pm.Normal("obs", mu.sum(), observed=y)
    pm.Deterministic("twice_mu", 2 * mu)   # change the multiplier to avoid full - graph cache leaking across runs

for call in range(3):
    start = perf_counter()
    pm.sample(
        model=model,
        nuts_sampler="nutpie",
        nuts={"adaptation": "low_rank"},
        compile_kwargs={"backend": "numba", "cache": True},
        tune=1, draws=1, chains=1, cores=1,
        random_seed=42,
        progressbar=False,
        compute_convergence_checks=False,
    )
    print(f"Call {call + 1}: {perf_counter() - start:.3f} s")

Before

Only 1 samples per chain. Reliable r-hat and ESS diagnostics require longer chains for accurate estimate.
NUTS[nutpie]: [mu]
Call 1: 2.626 s
Only 1 samples per chain. Reliable r-hat and ESS diagnostics require longer chains for accurate estimate.
NUTS[nutpie]: [mu]
Call 2: 2.122 s
Only 1 samples per chain. Reliable r-hat and ESS diagnostics require longer chains for accurate estimate.
NUTS[nutpie]: [mu]
Call 3: 2.099 s

After

Only 1 samples per chain. Reliable r-hat and ESS diagnostics require longer chains for accurate estimate.
NUTS[nutpie]: [mu]
Call 1: 1.308 s
Only 1 samples per chain. Reliable r-hat and ESS diagnostics require longer chains for accurate estimate.
NUTS[nutpie]: [mu]
Call 2: 0.275 s
Only 1 samples per chain. Reliable r-hat and ESS diagnostics require longer chains for accurate estimate.
NUTS[nutpie]: [mu]
Call 3: 0.268 s

I don't think the first time difference is real, sometimes goes one direction, sometimes the other. but the cache is clear to the naked eye