Numba is great, but you have to think about the order of your computations nonetheless.
What you would need is the “holy grail of all polyhedral optimizers 38” or so that would rewrite the order of your entire computation. There are projects like TVM or TensorFlow MLIR that are working on pushing some boundaries, but so far spending a few thoughts yourself is the best way to get these things faster.
Best regards
Thomas