@brandonwillard looking forward to see your progress on this, sounds very interesting as a use case for Numba. For you and others, there’s both this forum and numba/numba - Gitter for support and the issue tracker for issues, patches to Numba are also welcomed. Further, Numba has weekly meetings Weekly Public Meeting every Tuesday for 2021 for discussing more involved topics and research, feel free to suggest items for discussion.
As to the Numba backend itself, I’ve taken a look at the code in the backend for JAX. Assuming that you are prepared to relax the use of some of the perhaps more idiomatic python constructs, then just from a quick skim I’d say Numba has a reasonable chance of being able to compile a lot of the dispatched functions (as you’ve noted).
Do you also have plans to @jit compile the dispatched graph too so as to gain the benefits of IPO/loop fusion (you’ll have to excuse me probably not using the terminology common to Aesara)?