More setup singularity and update stencil

[ashwinvis]

FYI, @stepas-toliautas , this is a follow up since paths were changed in #113

[stepas-toliautas]

@ashwinvis : thanks for the updates!
Do I understand correctly that

• the CuPy version directly uses CUDA kernels from _cuda subfolder, and
• the Jax version has potential to use vendor-agnostic Python GPU kernels (as in SYCL for C++), but code still needs to be finished/ polished?

The latter would be really nice, but I won't have time to integrate Jax variant into existing codebase until tomorrow. Still, maybe a comment on the possibility during the lesson could be worth it.

[ashwinvis]

• CuPy version uses the kernel written inline as strings in the file, where we define the evolve kernel
• Jax needs a lot of polish certainly. It can be illuminating to show how it forces us to write code in a certain way. When it works, Jax is the most vendor agnostic of all.