[curvilinear_m]

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).

[killingtime74]

Numba doesn't support GPU though

[catshamando]

Yes it does: https://numba.readthedocs.io/en/stable/cuda/kernels.html

[killingtime74]

ah thanks, TIL

[the_svd_doctor]

It does (for NVIDIA at least)