Comment by ainch

if you want GPU-backed/compiled numpy code then JAX is also a good option.