I had good results doing an open box reimplementation. Gave qwen access to my old projects and it rebuilt it on JAX.

https://github.com/verdverm/pge-jax