Stable PyMC-to-numba API?

[WardBrian]

Hi all --

@bob-carpenter and I recently released 0.0.1 of walnutpie, which takes a lot of inspiration from this package and @aseyboldt's work on nuts-rs.

In particular, with an eye for compatibility, I modeled our callback type on the one nuts-rs uses. I'd love to allow users to efficiently plug in pymc models using numba like is done here, but it seems like it depends a lot on pymc internals.

Is there any chance that this kind of API could be exposed (either by nutpie or pymc itself) in a stable way we could depend on in other packages?

[ricardoV94]

Certainly. Whether it should be pymc or another package (need not be nutpie) depends on whether we need any rust (which we don't want to add to our pymc dist/pipeline). @aseyboldt will know

[WardBrian]

need not be nutpie

If it can't live in pymc directly, I wouldn't mind telling pymc users to also install nutpie to use walnutpie, but it would be extra nice if there was a small pymc-numba-bridge package that both could depend on without intertwining each other :)

[aseyboldt]

Moving that stuff to pymc would be great. There is no rust involved, just an annoying amount of pymc and numba internals. If someone finds a cleaner way to do this than what I managed to come up with, I'd also be more than happy.

[WardBrian]

I'm not terribly familiar with the internals of either package to be super helpful, but as long as we define a nice enough api boundary maybe someone can rewrite the internal pieces without it mattering too much

[WardBrian]

It seems very clear to me how to write a function compile_pymc_log_density(m: pm.Model) -> tuple[numba.CFunc, NDArray] based on what is in the repo here.

The feature where you can change data without recompiling means this probably can't be such a simple function, unfortunately. It can instead return a simple dataclass of those two items with the update function on that class, perhaps? (I don't expect walnutpie to necessarily expose that functionality, so whatever works for nutpie is fine by me)

If we can agree on an API I am happy to draft a PR which moves this functionality somewhere in pymc based on the existing code

[ricardoV94]

As a start if you can get code that works for you and still passes the nutpie CI that's already useful first sketch?

[aseyboldt]

If we are looking into this anyway, we could also check if we can find a way that is a little less complicated. What made this difficult is that we don't know the rank or types of the arrays we want to pass in userdata, so there needs to be some kind of generic programming that extracts the arrays from userdata on the numba side. But numba really wasn't designed to do generic programming, which is why I'm directly generating llvm code in `make_extraction_fn` through numbas extension API. An alternative to this would be to just put some strings together to generate python source code that can then be compiled through numba. This is what the pytensor does in a handful of places where similar restrictions come up. I don't like generating source code, but it might actually be cleaner (less unclean?) here. There is also a missing feature that I wanted to support for a while: deterministics that require rng (roughly the equivalent of random generated quantities in stan). This doesn't come up often in PyMC because things like posterior predictives etc are computed separately anyway, but is still would be nice to have from time to time.

…

On Wed, 5 Aug 2026, 22:52 Ricardo Vieira, ***@***.***> wrote: *ricardoV94* left a comment (pymc-devs/nutpie#336) <#336 (comment)> As a start if you can get code that works for you and still passes the nutpie CI that's already useful first sketch? — Reply to this email directly, view it on GitHub <#336?email_source=notifications&email_token=AAOLSHIOFXNTQNGKPTCA2M35IONBTA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTKMJZG4ZDAMZVG43KM4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJLDGN5XXIZLSL5RWY2LDNM#issuecomment-5197203576>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAOLSHKKYYABGKK2TNOK24T5IONBTAVCNFSNUABFKJSXA33TNF2G64TZHMZTANJUGU4TQMBQHNEXG43VMU5TKMBXGE4TMOBVGMZKC5QC> . Triage notifications, keep track of coding agent tasks and review pull requests on the go with GitHub Mobile for iOS <https://github.com/notifications/mobile/ios/AAOLSHK56EUAT5V5IBSZ6NT5IONBTA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTKMJZG4ZDAMZVG43KM4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJKTGN5XXIZLSL5UW64Y> and Android <https://github.com/notifications/mobile/android/AAOLSHIE54DNZUWZMME5CZD5IONBTA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTKMJZG4ZDAMZVG43KM4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJLTGN5XXIZLSL5QW4ZDSN5UWI>. Download it today! You are receiving this because you were mentioned.Message ID: ***@***.***>

[WardBrian]

If we are looking into this anyway, we could also check if we can find a way that is a little less complicated. What made this difficult is that we don't know the rank or types of the arrays we want to pass in userdata, so there needs to be some kind of generic programming that extracts the arrays from userdata on the numba sid

We must know them at the time that we're actually creating the logp pointer, right? It seems like types.Record.make_c_struct would let us pack arbitrary data into to the user data pointer in a nicer way, but I'm not positive I am following your point

[aseyboldt]

The tricky part was to turn the record into a tuple later. Claude came up with a solution that uses a bunch of intermediate functions:

def _make_extract_array(key, ndim, dtype):
    """Build a jitted function that reads one shared array out of user_data."""
    import numba
    from numba.np.unsafe.ndarray import to_fixed_tuple

    @numba.njit(inline="always")
    def extract_array(shared):
        data = numba.carray(
            address_as_void_pointer(shared["data"][key][()]),
            int(shared["size"][key][()]),
            dtype,
        )
        return data.reshape(to_fixed_tuple(shared["shape"][key], ndim))

    return extract_array


def _append_extractor(extract_rest, extract_one):
    @numba.njit(inline="always")
    def extract(user_data):
        return extract_rest(user_data) + (extract_one(user_data),)
 
    return extract
 
 
def make_extraction_chain(inner, arrays, indices, record_dtype, n_records):
    @numba.njit(inline="always")
    def extract_all(user_data):
        return ()
 
    for index in indices:
        val = arrays[index]
        extract_all = _append_extractor(
            extract_all, _make_extract_array(index, val.ndim, np.dtype(val.dtype))
        )
 
    @numba.njit(inline="always")
    def extract_shared(x, user_data_):
        user_data = numba.carray(user_data_, (n_records,), record_dtype)
        return inner(x, *extract_all(user_data))
 
    return extract_shared

But compile-time wise it is much better to just generate the code (numba doesn't do well with long chains of functions):

def make_extraction_flat(inner, arrays, indices, record_dtype, n_records):
   extractors = {
       f"_e{k}": _make_extract_array(i, arrays[i].ndim, np.dtype(arrays[i].dtype))
       for k, i in enumerate(indices)
   }
   args = ", ".join(f"_e{k}(user_data)" for k in range(len(indices)))
   src = (
       "def extract_shared(x, user_data_):\n"
       "    user_data = carray(user_data_, (n_records,), record_dtype)\n"
       f"    return inner(x, {args})\n"
   )
   ns = dict(
       extractors,
       inner=inner,
       carray=numba.carray,
       record_dtype=record_dtype,
       n_records=n_records,
   )
   exec(compile(src, "<extract_shared>", "exec"), ns)
   return numba.njit(inline="always")(ns["extract_shared"])

It also found the to_fixed_tuple function, which I had missed...

[WardBrian]

@aseyboldt ah, I see, because you need to splat that tuple into a function call.

String-code-gen is always a little spooky but I much prefer the second snippet to the first, for what it is worth.