I'm surprised to see pytorch and Jax mentioned as alternatives but not numba : https://github.com/numba/numba

I've recently had to implement a few kernels to lower the memory footprint and runtime of some pytorch function : it's been really nice because numba kernels have type hints support (as opposed to raw cupy kernels).