Survey of JIT and pre-compilation projects in Python

[benjeffery]

In #3135 we're planning to add some code to make it easier to work with numba to enable performant calculations on tree sequences without resorting to C code.

numba isn't the only tool of this sort so it seemed wise to do a quick survey of similar projects.

jax

Has a jit compiler for python functions, but focus of library is on ML/AI work as it does auto diff of those functions. Has own numpy-like module as jax.numpy. Requires different style code as jax arrays are immutable. Doesn't support jitted classes, analogous thing would be immutable dataclasses.

codon

Precompiles to native machine code an "almost Python" (ASCII strings, 64 bit int, unordered dicts, fixed-size tuples, some dynamic things don't work). Has it's own numpy and interops with existing Python code/libraries, so could call out to tskit to get arrays and then use them, or even call other tskit methods.

It also has a jit decorator, for just-in-time compilation, but doesn't support jitted classes.

mojo

Similar to codon but a much more different language that is "Python-like". Compiler is closed source, doesn't have own numpy, calls usual numpy through Python interop layer.

Others

Other similar things which don't seem worth looking at as speedups modest:
nuitka, PyPy

In summary, Codon looks very interesting!

[jeromekelleher]

So we don't necessarily have to use the jitclass layout with numba, it's just easier than passing around 15 arrays (or whatever it is) to the function. Maybe that would be less problematic overall, though, if we could more-or-less share the same code between (say) numba and codon?

[benjeffery]

Yes, was looking into this - with a version that keeps state in a set of arrays I have codon working. It seems to be about twice the runtime of numba, though.

[benjeffery]

Codon is different enough, and a bit slower, so I suggest we proceed with a numba-only module for now. I can't even use a tuple to keep the state arrays together as tuples have to be homogeneous.

[molpopgen]

Would it be useful to pick a "default" (numba, say) that is supported out of the box and other systems can be added later via pip's optional dependency syntax ([codon]) as needed?

[benjeffery]

We're not intending to add numba as a dependency as it's heavy and delayed with respect to python releases.
How I see the jit code module working is that it tries to import numba and gracefully messages that it is needed if absent.

[benjeffery]

Experimenting with numba:
Using a jitclass that has an iter_trees method that yields only costs 10% over a fully fused single function. I think this is worth it as the class will be cleanly reusable.