Yeah, that's definitely available in Python frameworks. (I was wrong to mention Numba - I think CuPy would be a better example: https://docs.cupy.dev/en/stable/overview.html)