The jupyter notebook contains a hardware accelerated finite volume simulation made using JAX.

You can read the write up about this in the article I wrote on Medium

The code is a port of the python tutorial made by Philip Mocz, https://github.com/pmocz/finitevolume-python/

Google Collab provides a GPU runtime, you can try this out using this link