Jitting like you mentioned is supported by the MAX graph API: https://docs.modular.com/max/tutorials/build-custom-ops. It could have a nicer syntax though to be more like Numba, I think you have an interesting idea there.