Update tokamax dependency to version 0.0.8

* This fixes an issue with immutabledict dependency missing in previous versions of Tokamax.
* JAX and Haiku have to be updated as well.

PiperOrigin-RevId: 854207701
Change-Id: I8c30856504f44e43e3a058e89b0866922ca56a5b

Author: Augustin-Zidek